From a8aafbdbfe71e1fbb0f86a3e133cfbba6a9686f5 Mon Sep 17 00:00:00 2001 From: Rohan Bierneni Date: Sun, 6 Sep 2026 04:30:06 +0000 Subject: [PATCH] [Stacked PR 1/5] Fix upstream pyink formatting and megablox sublane tiling - Reformat qwen3.py, qwen3_5.py, pyconfig_test.py, train_tests.py, and train_compile_test.py with pyink to conform with upstream CI style requirements. - Fix Megablox TPU v2 sublane tiling in pallas_mosaic_tpu_v2_gmm_kernel.py and pallas_mosaic_tpu_v2_tgmm_kernel.py to ensure correct tile dimension handling on TPU sublanes. --- .../pallas_mosaic_tpu_v2_gmm_kernel.py | 2 +- .../pallas_mosaic_tpu_v2_tgmm_kernel.py | 3 + src/maxtext/models/qwen3.py | 65 ++----- src/maxtext/models/qwen3_5.py | 36 +--- tests/integration/train_tests.py | 19 +- tests/unit/pyconfig_test.py | 8 +- tests/unit/train_compile_test.py | 182 +++++++++--------- 7 files changed, 126 insertions(+), 189 deletions(-) diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 1826ef576b..4f3c70f19b 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -1090,7 +1090,7 @@ def validate_inputs( assert group_offset.shape == (1,) - size_lhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype) + size_lhs_sublane = max(pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype), 16) size_lhs_sublane = min(size_lhs_sublane, size_m) if fuse_act is not None: num_lanes = pltpu.get_tpu_info().num_lanes diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py index a5f3f6338d..2bf524a5e4 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py @@ -221,6 +221,9 @@ def make_tgmm_configs( size_lhs_sublane = min(size_lhs_sublane, size_m) size_rhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(rhs.dtype) size_rhs_sublane = min(size_rhs_sublane, size_m) + common_sublane = min(size_lhs_sublane, size_rhs_sublane) + size_lhs_sublane = common_sublane + size_rhs_sublane = common_sublane assert size_lhs_sublane == size_rhs_sublane, ( f"size_lhs_sublane should be the same as size_rhs_sublane {lhs.dtype=}," f" {rhs.dtype=}" ) diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 1a04df4c7c..24da5a7034 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -13,6 +13,7 @@ # limitations under the License. """Qwen3 family of model decoder layers.""" + # pylint: disable=arguments-differ # pylint: disable=no-name-in-module @@ -624,9 +625,7 @@ def _explicit_activation_shardings(self, batch: int): cp_len = LENGTH if gdn_context_axes(self.config) else None def _sharding(logical_axes): - pspec = logical_to_mesh_axes( - logical_axes, mesh=self.mesh, rules=logical_rules - ) + pspec = logical_to_mesh_axes(logical_axes, mesh=self.mesh, rules=logical_rules) # Training microbatches can be smaller than the physical batch partition. # Only dim 0 is inspected, so the trailing sizes are placeholders. shape = (batch,) + (1,) * (len(logical_axes) - 1) @@ -654,9 +653,7 @@ def __call__( # hidden_states: (B, S, E) cfg = self.config batch, seq_len, _ = hidden_states.shape - flat_sharding, head_sharding, state_sharding = ( - self._explicit_activation_shardings(batch) - ) + flat_sharding, head_sharding, state_sharding = self._explicit_activation_shardings(batch) active_cache = kv_cache if kv_cache is not None else self.cache @@ -759,13 +756,9 @@ def __call__( b_raw, a_raw = jnp.split(mixed_ba, split_indices_ba, axis=3) # b: (B, S, H_v) - b = jnp.reshape( - b_raw, (batch, seq_len, self.num_v_heads), out_sharding=flat_sharding - ) + b = jnp.reshape(b_raw, (batch, seq_len, self.num_v_heads), out_sharding=flat_sharding) # a: (B, S, H_v) - a = jnp.reshape( - a_raw, (batch, seq_len, self.num_v_heads), out_sharding=flat_sharding - ) + a = jnp.reshape(a_raw, (batch, seq_len, self.num_v_heads), out_sharding=flat_sharding) if use_paged_state: # ========================================================================= @@ -1051,9 +1044,7 @@ def extract_state(c_in, v_len): value = jax.sharding.reshard(value, qkv_pspec) g = jax.sharding.reshard(g, g_beta_pspec) beta = jax.sharding.reshard(beta, g_beta_pspec) - recurrent_state_arg = jax.sharding.reshard( - recurrent_state_arg, state_pspec - ) + recurrent_state_arg = jax.sharding.reshard(recurrent_state_arg, state_pspec) @functools.partial( jax.shard_map, @@ -1127,15 +1118,11 @@ def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): # The normalization and gating is applied per-head on the value dimension. # Apply the norm and gate. Output shape: (B, S, H_v, D_v) - gated_output_reshaped = self.norm( - core_attn_out, z, out_sharding=head_sharding - ) + gated_output_reshaped = self.norm(core_attn_out, z, out_sharding=head_sharding) # Reshape back to a single feature dimension for the final projection. # Shape from (B, S, H_v, D_v) -> (B, S, value_dim) - gated_output = jnp.reshape( - gated_output_reshaped, (batch, seq_len, -1), out_sharding=flat_sharding - ) + gated_output = jnp.reshape(gated_output_reshaped, (batch, seq_len, -1), out_sharding=flat_sharding) # Final output shape: (B, S, E) output = self.out_proj(gated_output, out_sharding=out_sharding) @@ -1692,9 +1679,7 @@ def __init__( # Physical shardings used to pin sublayer outputs under ShardMode.EXPLICIT. In # ShardMode.AUTO the callees ignore these and let GSPMD infer the layout. if cfg.shard_mode == ShardMode.EXPLICIT: - self.out_sharding = create_sharding( - mesh, self.activation_axis_names, rules=get_logical_axis_rules() - ) + self.out_sharding = create_sharding(mesh, self.activation_axis_names, rules=get_logical_axis_rules()) self.mlp_intermediate_sharding = create_sharding( mesh, self.mlp_activation_axis_names, rules=get_logical_axis_rules() ) @@ -1777,15 +1762,11 @@ def __call__( # First LayerNorm, applied before the attention block. hidden_states = self.input_layernorm(inputs, out_sharding=self.out_sharding) - hidden_states = self._maybe_shard_with_logical( - hidden_states, self.activation_axis_names - ) + hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) # Conditionally apply either the Linear Attention or Full Attention block. if isinstance(self.attention, Qwen3NextFullAttention): - attention_output, new_kv_cache = cast( - Qwen3NextFullAttention, self.attention - )( + attention_output, new_kv_cache = cast(Qwen3NextFullAttention, self.attention)( hidden_states, decoder_segment_ids, decoder_positions, @@ -1796,9 +1777,7 @@ def __call__( out_sharding=self.out_sharding, ) else: - attention_output, new_kv_cache = cast( - Qwen3NextGatedDeltaNet, self.attention - )( + attention_output, new_kv_cache = cast(Qwen3NextGatedDeltaNet, self.attention)( hidden_states, model_mode=model_mode, kv_cache=kv_cache, @@ -1808,24 +1787,16 @@ def __call__( ) # First residual connection after attention - attention_output = self._maybe_shard_with_logical( - attention_output, self.activation_axis_names - ) + attention_output = self._maybe_shard_with_logical(attention_output, self.activation_axis_names) hidden_states = residual + attention_output - hidden_states = self._maybe_shard_with_logical( - hidden_states, self.activation_axis_names - ) + hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) # Prepare for the MoE block by capturing the new residual residual = hidden_states # Second LayerNorm, applied before the MoE block. - hidden_states = self.post_attention_layernorm( - hidden_states, out_sharding=self.out_sharding - ) - hidden_states = self._maybe_shard_with_logical( - hidden_states, self.activation_axis_names - ) + hidden_states = self.post_attention_layernorm(hidden_states, out_sharding=self.out_sharding) + hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) # Instantiate and call our `Qwen3NextSparseMoeBlock`. mlp_output, load_balance_loss = self.mlp( @@ -1841,9 +1812,7 @@ def __call__( self.moe_lb_loss = nnx.Intermediate(load_balance_loss) # Final residual connection (after the MoE block) - mlp_output = self._maybe_shard_with_logical( - mlp_output, self.activation_axis_names - ) + mlp_output = self._maybe_shard_with_logical(mlp_output, self.activation_axis_names) layer_output = residual + mlp_output layer_output = self._maybe_shard_with_logical( layer_output, diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py index b80ef0edf5..e54c82c730 100644 --- a/src/maxtext/models/qwen3_5.py +++ b/src/maxtext/models/qwen3_5.py @@ -144,9 +144,7 @@ def __init__( # Physical shardings used to pin sublayer outputs under ShardMode.EXPLICIT. In # ShardMode.AUTO the callees ignore these and let GSPMD infer the layout. if cfg.shard_mode == ShardMode.EXPLICIT: - self.out_sharding = create_sharding( - mesh, self.activation_axis_names, rules=get_logical_axis_rules() - ) + self.out_sharding = create_sharding(mesh, self.activation_axis_names, rules=get_logical_axis_rules()) self.mlp_intermediate_sharding = create_sharding( mesh, self.mlp_activation_axis_names, rules=get_logical_axis_rules() ) @@ -226,15 +224,11 @@ def __call__( # First LayerNorm, applied before the attention block. hidden_states = self.input_layernorm(inputs, out_sharding=self.out_sharding) - hidden_states = self._maybe_shard_with_logical( - hidden_states, self.activation_axis_names - ) + hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) # Conditionally apply either the Linear Attention or Full Attention block. if isinstance(self.attention, Qwen3_5FullAttention): - attention_output, new_kv_cache = cast( - Qwen3_5FullAttention, self.attention - )( + attention_output, new_kv_cache = cast(Qwen3_5FullAttention, self.attention)( hidden_states, decoder_segment_ids, decoder_positions, @@ -245,9 +239,7 @@ def __call__( out_sharding=self.out_sharding, ) else: - attention_output, new_kv_cache = cast( - Qwen3_5GatedDeltaNet, self.attention - )( + attention_output, new_kv_cache = cast(Qwen3_5GatedDeltaNet, self.attention)( hidden_states, model_mode=model_mode, kv_cache=kv_cache, @@ -257,24 +249,16 @@ def __call__( ) # First residual connection after attention - attention_output = self._maybe_shard_with_logical( - attention_output, self.activation_axis_names - ) + attention_output = self._maybe_shard_with_logical(attention_output, self.activation_axis_names) hidden_states = residual + attention_output - hidden_states = self._maybe_shard_with_logical( - hidden_states, self.activation_axis_names - ) + hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) # Prepare for the MoE block by capturing the new residual residual = hidden_states # Second LayerNorm, applied before the MoE block. - hidden_states = self.post_attention_layernorm( - hidden_states, out_sharding=self.out_sharding - ) - hidden_states = self._maybe_shard_with_logical( - hidden_states, self.activation_axis_names - ) + hidden_states = self.post_attention_layernorm(hidden_states, out_sharding=self.out_sharding) + hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) # Instantiate and call our `Qwen3_5SparseMoEBlock`. mlp_output, load_balance_loss = self.mlp( @@ -291,9 +275,7 @@ def __call__( self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss) # Final residual connection (after the MoE block) - mlp_output = self._maybe_shard_with_logical( - mlp_output, self.activation_axis_names - ) + mlp_output = self._maybe_shard_with_logical(mlp_output, self.activation_axis_names) layer_output = residual + mlp_output layer_output = self._maybe_shard_with_logical( layer_output, diff --git a/tests/integration/train_tests.py b/tests/integration/train_tests.py index a83016709f..711efcd5b5 100644 --- a/tests/integration/train_tests.py +++ b/tests/integration/train_tests.py @@ -222,9 +222,7 @@ class TrainTests(unittest.TestCase): # The Qwen3.5 model configs default to a HuggingFace tokenizer that is not # vendored in the repo; use the checked-in tiktoken asset instead. "tokenizer_type=tiktoken", - ( - rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}" - ), + (rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}"), ] # Same sublayers as Qwen3.5, wired together by Qwen3NextScannableBlock rather than a @@ -1072,26 +1070,20 @@ def test_tpu_qwen3_hybrid_explicit_sharding_matches_auto(self): for decoder_block, model_overrides in self._QWEN3_HYBRID_MODELS.items(): with self.subTest(decoder_block=decoder_block): args = parallelism[decoder_block] - auto_losses = self._losses( - f"{decoder_block}_auto", model_overrides, args + ["shard_mode=auto"] - ) + auto_losses = self._losses(f"{decoder_block}_auto", model_overrides, args + ["shard_mode=auto"]) explicit_losses = self._losses( f"{decoder_block}_explicit", model_overrides, args + ["shard_mode=explicit"], ) print(f"[{decoder_block}] auto losses: {auto_losses}", flush=True) - print( - f"[{decoder_block}] explicit losses: {explicit_losses}", flush=True - ) + print(f"[{decoder_block}] explicit losses: {explicit_losses}", flush=True) self.assertTrue(auto_losses, "auto run produced no metrics") # `activation_batch` carries the expert axis, so pinning it reassociates the # backward reductions: the forward pass is bit-for-bit and the drift only appears # once gradients flow. Over 20 steps it stays below 4e-5 relative and changes # sign, i.e. it is float noise rather than the two runs pulling apart. - np.testing.assert_allclose( - explicit_losses, auto_losses, rtol=1e-4, atol=0.0 - ) + np.testing.assert_allclose(explicit_losses, auto_losses, rtol=1e-4, atol=0.0) @pytest.mark.integration_test @pytest.mark.tpu_only @@ -1127,8 +1119,7 @@ def test_tpu_qwen3_hybrid_zero1_gradient_accumulation(self): sharded = self._losses( f"{decoder_block}_ga_zero1", model_overrides, - zero1_ga - + ["shard_mode=explicit", "shard_optimizer_over_data=True"], + zero1_ga + ["shard_mode=explicit", "shard_optimizer_over_data=True"], ) print(f"[{decoder_block}] auto + GA losses: {baseline}", flush=True) print( diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py index 7e272fc581..d870026e35 100644 --- a/tests/unit/pyconfig_test.py +++ b/tests/unit/pyconfig_test.py @@ -285,14 +285,10 @@ def initialize(decoder_block=decoder_block, **kwargs): with self.assertRaisesRegex(Exception, "requires `sparse_matmul=True`"): initialize(sparse_matmul=False, megablox=False) - with self.assertRaisesRegex( - Exception, "does not support context parallelism" - ): + with self.assertRaisesRegex(Exception, "does not support context parallelism"): initialize(ici_context_parallelism=2) - with self.assertRaisesRegex( - Exception, "does not support context parallelism" - ): + with self.assertRaisesRegex(Exception, "does not support context parallelism"): initialize(ici_context_usp_ulysses_parallelism=2) def test_explicit_sharding_mistral_decoder_support(self): diff --git a/tests/unit/train_compile_test.py b/tests/unit/train_compile_test.py index 70653e4ec7..f22b90ea81 100644 --- a/tests/unit/train_compile_test.py +++ b/tests/unit/train_compile_test.py @@ -1149,53 +1149,53 @@ def test_qwen3_next_explicit_sharding(self): the layout it was handed. """ - compiled_trainstep_file = os.path.join( - gettempdir(), "test_qwen3_next_explicit_sharding.pickle" + compiled_trainstep_file = os.path.join(gettempdir(), "test_qwen3_next_explicit_sharding.pickle") + train_compile_main( + ( + "", + get_test_config_path(), + f"compiled_trainstep_file={compiled_trainstep_file}", + "compile_topology=v5p-512", + "compile_topology_num_slices=1", + "model_name=qwen3-next-80b-a3b", + "per_device_batch_size=1.0", + "max_target_length=1024", + "ici_fsdp_parallelism=32", + "ici_expert_parallelism=8", + "sparse_matmul=True", + "megablox=True", + "attention=flash", + "use_tokamax_splash=True", + "shard_mode=explicit", + ) ) - train_compile_main(( - "", - get_test_config_path(), - f"compiled_trainstep_file={compiled_trainstep_file}", - "compile_topology=v5p-512", - "compile_topology_num_slices=1", - "model_name=qwen3-next-80b-a3b", - "per_device_batch_size=1.0", - "max_target_length=1024", - "ici_fsdp_parallelism=32", - "ici_expert_parallelism=8", - "sparse_matmul=True", - "megablox=True", - "attention=flash", - "use_tokamax_splash=True", - "shard_mode=explicit", - )) def test_qwen3_next_explicit_sharding_zero1(self): """AOT test for qwen3-next under explicit sharding with ZeRO-1 and gradient accumulation.""" - compiled_trainstep_file = os.path.join( - gettempdir(), "test_qwen3_next_explicit_sharding_zero1.pickle" + compiled_trainstep_file = os.path.join(gettempdir(), "test_qwen3_next_explicit_sharding_zero1.pickle") + train_compile_main( + ( + "", + get_test_config_path(), + f"compiled_trainstep_file={compiled_trainstep_file}", + "compile_topology=v5p-256", + "compile_topology_num_slices=1", + "model_name=qwen3-next-80b-a3b", + "override_model_config=True", + "base_num_decoder_layers=4", + "per_device_batch_size=1.0", + "max_target_length=1024", + "sparse_matmul=True", + "megablox=True", + "attention=flash", + "use_tokamax_splash=True", + "shard_mode=explicit", + "ici_data_parallelism=-1", + "ici_fsdp_parallelism=1", + "gradient_accumulation_steps=4", + "shard_optimizer_over_data=True", + ) ) - train_compile_main(( - "", - get_test_config_path(), - f"compiled_trainstep_file={compiled_trainstep_file}", - "compile_topology=v5p-256", - "compile_topology_num_slices=1", - "model_name=qwen3-next-80b-a3b", - "override_model_config=True", - "base_num_decoder_layers=4", - "per_device_batch_size=1.0", - "max_target_length=1024", - "sparse_matmul=True", - "megablox=True", - "attention=flash", - "use_tokamax_splash=True", - "shard_mode=explicit", - "ici_data_parallelism=-1", - "ici_fsdp_parallelism=1", - "gradient_accumulation_steps=4", - "shard_optimizer_over_data=True", - )) def test_qwen3_5_explicit_sharding(self): """AOT test for qwen3-5 under explicit sharding, at FSDP 32 x expert 8. @@ -1206,31 +1206,29 @@ def test_qwen3_5_explicit_sharding(self): silently costing a collective at a scale a real test cannot reach. """ - compiled_trainstep_file = os.path.join( - gettempdir(), "test_qwen3_5_explicit_sharding.pickle" - ) - train_compile_main(( - "", - get_test_config_path(), - f"compiled_trainstep_file={compiled_trainstep_file}", - "compile_topology=v5p-512", - "compile_topology_num_slices=1", - "model_name=qwen3.5-397b-a17b", - "per_device_batch_size=1.0", - "max_target_length=1024", - "ici_fsdp_parallelism=32", - "ici_expert_parallelism=8", - "sparse_matmul=True", - "megablox=True", - "attention=flash", - "use_tokamax_splash=True", - "shard_mode=explicit", - # Qwen3.5 defaults to a HuggingFace tokenizer that is not vendored. - "tokenizer_type=tiktoken", + compiled_trainstep_file = os.path.join(gettempdir(), "test_qwen3_5_explicit_sharding.pickle") + train_compile_main( ( - f"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}" - ), - )) + "", + get_test_config_path(), + f"compiled_trainstep_file={compiled_trainstep_file}", + "compile_topology=v5p-512", + "compile_topology_num_slices=1", + "model_name=qwen3.5-397b-a17b", + "per_device_batch_size=1.0", + "max_target_length=1024", + "ici_fsdp_parallelism=32", + "ici_expert_parallelism=8", + "sparse_matmul=True", + "megablox=True", + "attention=flash", + "use_tokamax_splash=True", + "shard_mode=explicit", + # Qwen3.5 defaults to a HuggingFace tokenizer that is not vendored. + "tokenizer_type=tiktoken", + (f"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}"), + ) + ) def test_qwen3_5_explicit_sharding_zero1(self): """AOT test for qwen3-5 under explicit sharding with ZeRO-1 and gradient accumulation. @@ -1245,34 +1243,32 @@ def test_qwen3_5_explicit_sharding_zero1(self): variants while staying small enough to hold data-parallel replicas of the parameters. """ - compiled_trainstep_file = os.path.join( - gettempdir(), "test_qwen3_5_explicit_sharding_zero1.pickle" - ) - train_compile_main(( - "", - get_test_config_path(), - f"compiled_trainstep_file={compiled_trainstep_file}", - "compile_topology=v5p-256", - "compile_topology_num_slices=1", - "model_name=qwen3.5-35b-a3b", - "override_model_config=True", - "base_num_decoder_layers=4", - "per_device_batch_size=1.0", - "max_target_length=1024", - "sparse_matmul=True", - "megablox=True", - "attention=flash", - "use_tokamax_splash=True", - "shard_mode=explicit", - "ici_data_parallelism=-1", - "ici_fsdp_parallelism=1", - "gradient_accumulation_steps=4", - "shard_optimizer_over_data=True", - "tokenizer_type=tiktoken", + compiled_trainstep_file = os.path.join(gettempdir(), "test_qwen3_5_explicit_sharding_zero1.pickle") + train_compile_main( ( - f"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}" - ), - )) + "", + get_test_config_path(), + f"compiled_trainstep_file={compiled_trainstep_file}", + "compile_topology=v5p-256", + "compile_topology_num_slices=1", + "model_name=qwen3.5-35b-a3b", + "override_model_config=True", + "base_num_decoder_layers=4", + "per_device_batch_size=1.0", + "max_target_length=1024", + "sparse_matmul=True", + "megablox=True", + "attention=flash", + "use_tokamax_splash=True", + "shard_mode=explicit", + "ici_data_parallelism=-1", + "ici_fsdp_parallelism=1", + "gradient_accumulation_steps=4", + "shard_optimizer_over_data=True", + "tokenizer_type=tiktoken", + (f"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}"), + ) + ) def test_serialization_and_deserialization_formats(self): """Tests that our custom binary save/load functions work securely and legacy fallback triggers warning."""