From 60591b3d629c8da8e97433b21fa6f576ac782363 Mon Sep 17 00:00:00 2001 From: ethannnnnn Date: Wed, 22 Jul 2026 22:29:00 -0700 Subject: [PATCH 1/6] [maxtext] Add generalized block-causal attention Add block_diffusion as a default-off attention type with one namespaced block-size setting. Resolve it in the shared attention layer so ordinary global layers need no model-name dispatch while explicit specialized attention remains intact. Implement matching dense, Splash, Tokamax, packed-sequence, and load-balanced context-parallel masks. Preserve the autoregressive path and support partial final blocks. Tests: attention 42 passed/37 skipped; config 18 passed plus 21 subtests; Tokamax 3 passed; pyink, yamllint, git diff --check. --- src/maxtext/common/common_types.py | 1 + src/maxtext/configs/base.yml | 3 +- src/maxtext/configs/types.py | 14 +- src/maxtext/layers/attention_op.py | 124 +++++++++++++- src/maxtext/layers/attentions.py | 8 +- tests/unit/attention_test.py | 154 ++++++++++++++++++ tests/unit/configs_value_test.py | 48 ++++++ .../tokamax_splash_attention_mask_test.py | 17 ++ 8 files changed, 360 insertions(+), 9 deletions(-) diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index 9dd4aff3dc..8ed18bfd12 100644 --- a/src/maxtext/common/common_types.py +++ b/src/maxtext/common/common_types.py @@ -135,6 +135,7 @@ class AttentionType(enum.Enum): MLA = "mla" COMPRESSED = "compressed" FULL = "full" + BLOCK_DIFFUSION = "block_diffusion" class ShardMode(enum.Enum): diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 55925ffd58..4d0bf560e8 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -380,12 +380,13 @@ param_scan_axis: 1 # The attention parameter dictates the specific algorithm/methodology used to compute the attention scores # The attention_type parameter determines the variants of attention, e.g. global or local_sliding attention: 'autoselected' # Supported attention: autoselected, dot_product, flash, cudnn_flash_te -attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla +attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla, full, compressed, block_diffusion share_kv_projections: false # Note: Not compatible with attention_type='mla' attention_bias: false # If true, adds a learnable bias to the query, key, and value projections attention_sink: false sliding_window_size: 0 chunk_attn_window_size: 0 +block_diffusion_block_size: 32 # Size of each bidirectional block when attention_type='block_diffusion' attn_logits_soft_cap: 0.0 final_logits_soft_cap: 0.0 z_loss_multiplier: 0.0 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 03640ec758..8e42109894 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -572,7 +572,7 @@ class Attention(BaseModel): "autoselected", description="The attention algorithm to use (dot_product, flash, cudnn_flash_te, vllm_rpa, vllm_batched_rpa, etc).", ) - attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed"] = Field( + attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed", "block_diffusion"] = Field( "global", description="The variant of attention to use." ) share_kv_projections: bool = Field( @@ -619,6 +619,10 @@ class Attention(BaseModel): ) sliding_window_size: NonNegativeInt = Field(0, description="The size of the sliding window for local attention.") chunk_attn_window_size: NonNegativeInt = Field(0, description="The window size for chunked attention.") + block_diffusion_block_size: PositiveInt = Field( + 32, + description="The size of each bidirectional block when attention_type is 'block_diffusion'.", + ) attn_logits_soft_cap: None | NonNegativeFloat = Field( None, description="Soft-cap value for attention logits. None means no cap." ) @@ -3283,6 +3287,14 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de not isinstance(self.sliding_window_size, int) or self.sliding_window_size <= 0 ): raise ValueError("`sliding_window_size` must be an integer > 0 for 'local_sliding' attention.") + if self.attention_type == AttentionType.BLOCK_DIFFUSION.value: + if self.attention not in ("autoselected", "dot_product", "flash"): + raise ValueError("Block-diffusion attention is supported only by dot_product attention and TPU Splash attention.") + if self.attention in ("autoselected", "flash") and self.hardware != "tpu": + raise ValueError( + "Block-diffusion attention with attention='autoselected' or attention='flash' requires hardware='tpu'; " + "use attention='dot_product' on other hardware." + ) if self.quantize_kvcache and not self.kv_quant_axis: raise ValueError("`kv_quant_axis` cannot be empty when quantize_kvcache is True.") if ( diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 8e6c96a72e..af48802cc9 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -84,6 +84,16 @@ dynamic_vector_slice_in_dim = jax.vmap(lax.dynamic_slice_in_dim, in_axes=(None, 0, None, None)) +def _resolve_attention_type(config: Config, attention_type: AttentionType | str | None) -> AttentionType: + configured_attention_type = AttentionType(getattr(config, "attention_type", AttentionType.GLOBAL.value)) + if attention_type is None: + return configured_attention_type + resolved_attention_type = AttentionType(attention_type) + if configured_attention_type == AttentionType.BLOCK_DIFFUSION and resolved_attention_type == AttentionType.GLOBAL: + return configured_attention_type + return resolved_attention_type + + def validate_compute_axis_order(s: AxisIdxes) -> None: valid_compute_axis_order = ((0, 1, 2, 3), (0, 2, 1, 3)) if s not in valid_compute_axis_order: # currently supported compute_axis_order @@ -204,6 +214,50 @@ def __hash__(self): ) +class BlockCausalMask(splash_attention_mask._ComputableMask): # pylint: disable=protected-access,abstract-method + """Lazy mask with bidirectional attention within causal blocks.""" + + block_size: int + + def __init__( + self, + shape: tuple[int, int], + block_size: int, + shard_count: int = 1, + ): + if block_size <= 0: + raise ValueError("block_size must be positive") + self.block_size = block_size + + def block_causal_mask_function(q_ids, kv_ids): + return (q_ids // self.block_size) >= (kv_ids // self.block_size) + + super().__init__( + shape=shape, + mask_function=block_causal_mask_function, + shard_count=shard_count, + ) + + def __eq__(self, other: object): + if not isinstance(other, type(self)): + return NotImplemented + return ( + self.shape == other.shape + and self.block_size == other.block_size + and np.array_equal(self.q_sequence, other.q_sequence) + ) + + def __hash__(self): + return hash( + ( + type(self), + self.shape, + self.block_size, + self.q_sequence.tobytes() if self.q_sequence is not None else None, + ) + ) + + def _generate_chunk_attention_mask(mask_shape: tuple[int, int], chunk_size: int, q_offset: int = 0) -> jax.Array: """Generates an explicit boolean mask for chunked causal attention. @@ -234,6 +288,16 @@ def _generate_chunk_attention_mask(mask_shape: tuple[int, int], chunk_size: int, return chunk_mask +def _generate_block_causal_attention_mask(mask_shape: tuple[int, int], block_size: int, q_offset: int = 0) -> jax.Array: + """Generates a block-causal mask, bidirectional within each block.""" + if block_size <= 0: + raise ValueError("block_size must be positive") + + row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0) + q_offset + col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1) + return (row_ids // block_size) >= (col_ids // block_size) + + def _make_block_mask_indices(bidirectional_mask): """Creates block mask identifying segments based on a bidirectional mask. @@ -479,7 +543,13 @@ def __init__( self.dtype = dtype self.quant = quant self.kv_quant = kv_quant - self.attention_type = attention_type + self.attention_type = _resolve_attention_type(self.config, attention_type) + self.block_diffusion_block_size = getattr(self.config, "block_diffusion_block_size", None) + if self.attention_type == AttentionType.BLOCK_DIFFUSION: + if self.block_diffusion_block_size is None or self.block_diffusion_block_size <= 0: + raise ValueError("block_diffusion_block_size must be positive for block-diffusion attention") + if self.attention_kernel not in ("autoselected", "dot_product", "flash"): + raise ValueError("Block-diffusion attention is supported only by dot_product attention and TPU Splash attention.") # Block sizes are only used by TPU splash attention kernels. Exclude non-splash kernels if self.attention_kernel not in ( "dot_product", @@ -782,6 +852,7 @@ def generate_attention_mask( if model_mode != MODEL_MODE_AUTOREGRESSIVE and self.attention_type not in ( AttentionType.FULL, AttentionType.COMPRESSED, + AttentionType.BLOCK_DIFFUSION, ): if use_segment_positions: causal_mask = (position_col_ids <= position_row_ids)[:, None, None, :, :] @@ -861,6 +932,19 @@ def generate_attention_mask( ) output_mask = chunk_mask * output_mask + elif self.attention_type == AttentionType.BLOCK_DIFFUSION and model_mode != MODEL_MODE_AUTOREGRESSIVE: + if use_segment_positions: + block_mask = ( + (position_row_ids // self.block_diffusion_block_size) >= (position_col_ids // self.block_diffusion_block_size) + )[:, None, None, :, :] + else: + block_mask = _generate_block_causal_attention_mask( + mask_shape=(q_seq_len, kv_seq_len), + block_size=self.block_diffusion_block_size, + q_offset=next_pos, + )[None, None, None, :, :] + output_mask = block_mask if output_mask is None else jnp.logical_and(output_mask, block_mask) + if bidirectional_mask is not None: image_mask = _make_bidirectional_block_mask(bidirectional_mask) output_mask = output_mask | image_mask[:, None, None, ...] @@ -1126,6 +1210,11 @@ def apply_attention( return out, None, None else: + if self.attention_type == AttentionType.BLOCK_DIFFUSION: + raise ValueError( + "Block-diffusion flash attention is supported only by TPU Splash; " + "use attention='dot_product' on other hardware." + ) if model_mode == MODEL_MODE_AUTOREGRESSIVE: # fallback to dot_product as pallas gpu flash attention doesn't support decode stage return self.apply_attention_dot( @@ -1430,13 +1519,24 @@ def create_sa_config(config, query, key, attn_logits_soft_cap): sa_config = create_sa_config(self.config, query, key, attn_logits_soft_cap) mask_shape = (query.shape[2], key.shape[2]) # (q_seq_len, kv_seq_len) mask_module = tokamax_splash_mask if self.config.use_tokamax_splash else splash_attention_mask + use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel if self.attention_type == AttentionType.FULL: mask = mask_module.FullMask(mask_shape) + elif self.attention_type == AttentionType.BLOCK_DIFFUSION: + mask_type = LoadBalancedBlockCausalMask if use_load_balanced_cp else BlockCausalMask + mask_kwargs = {"cp_size": cp_size} if use_load_balanced_cp else {} + mask = mask_type( + shape=mask_shape, + block_size=self.block_diffusion_block_size, + **mask_kwargs, + ) else: mask = mask_module.CausalMask(shape=mask_shape) - use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel - if use_load_balanced_cp and self.attention_type != AttentionType.FULL: + if use_load_balanced_cp and self.attention_type not in ( + AttentionType.FULL, + AttentionType.BLOCK_DIFFUSION, + ): mask = LoadBalancedCausalMask(shape=mask_shape, cp_size=cp_size) # Apply local masking if local sliding attention is enabled. @@ -2447,3 +2547,21 @@ def __init__( shard_count=shard_count, ) self.q_sequence = _load_balanced_q_sequence(shape, cp_size) + + +class LoadBalancedBlockCausalMask(BlockCausalMask): # pylint: disable=abstract-method + """Lazy block-causal mask with load-balanced query positions.""" + + def __init__( + self, + shape: tuple[int, int], + block_size: int, + cp_size: int, + shard_count: int = 1, + ): + super().__init__( + shape=shape, + block_size=block_size, + shard_count=shard_count, + ) + self.q_sequence = _load_balanced_q_sequence(shape, cp_size) diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index 5a65b35ba4..32f233cbeb 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -51,7 +51,7 @@ AttentionType, ) from maxtext.layers import nnx_wrappers -from maxtext.layers.attention_op import AttentionOp +from maxtext.layers.attention_op import AttentionOp, _resolve_attention_type from maxtext.layers.embeddings import ( LLaMARotaryEmbedding, LlamaVisionRotaryEmbedding, @@ -120,7 +120,7 @@ def attention_as_linen( float32_logits: bool = False, # cast logits in float32 for stability. quant: Optional[Quant] = None, kv_quant: Optional[KVQuant] = None, - attention_type: AttentionType = AttentionType.GLOBAL, # Default to global attention + attention_type: AttentionType | None = None, attn_logits_soft_cap: float | None = None, sliding_window_size: int | None = None, use_ragged_attention: bool = False, @@ -277,7 +277,7 @@ def __init__( float32_logits: bool = False, # cast logits in float32 for stability. quant: Optional[Quant] = None, kv_quant: Optional[KVQuant] = None, - attention_type: AttentionType = AttentionType.GLOBAL, # Default to global attention + attention_type: AttentionType | None = None, attn_logits_soft_cap: float | None = None, sliding_window_size: int | None = None, use_ragged_attention: bool = False, @@ -388,7 +388,7 @@ def __init__( self.float32_logits = float32_logits self.quant = quant self.kv_quant = kv_quant - self.attention_type = attention_type + self.attention_type = _resolve_attention_type(self.config, attention_type) self.attn_logits_soft_cap = attn_logits_soft_cap self.sliding_window_size = sliding_window_size self.use_ragged_attention = use_ragged_attention diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 20522bbd3b..0c7aa03005 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -44,9 +44,12 @@ from maxtext.layers import attention_op from maxtext.layers.attention_op import ( AttentionOp, + BlockCausalMask, ChunkedCausalMask, + _generate_block_causal_attention_mask, _generate_chunk_attention_mask, _make_bidirectional_block_mask, + _resolve_attention_type, ) from maxtext.layers.attentions import Attention from maxtext.layers import embeddings @@ -330,6 +333,143 @@ def test_value_error_on_zero_chunk_size(self): _generate_chunk_attention_mask(mask_shape=(4, 4), chunk_size=0) +class BlockCausalMaskTest(unittest.TestCase): + """Tests the shared dense and Splash block-causal masks.""" + + def _make_op(self, sequence_length, *, attention_type=AttentionType.BLOCK_DIFFUSION): + """Builds a minimal dot-product attention operator.""" + config = types.SimpleNamespace( + block_diffusion_block_size=4, + context_parallel_load_balance=False, + context_sharding="context", + ) + mesh = types.SimpleNamespace(shape={}) + kwargs = {} + if attention_type is not None: + kwargs["attention_type"] = attention_type + return AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=sequence_length, + mesh=mesh, + attention_kernel="dot_product", + **kwargs, + ) + + def test_dense_and_splash_masks_match(self): + sequence_length = 10 + block_size = 4 + query_positions = np.arange(sequence_length)[:, None] + key_positions = np.arange(sequence_length)[None, :] + expected = query_positions // block_size >= key_positions // block_size + + splash_mask = BlockCausalMask((sequence_length, sequence_length), block_size) + dense_mask = _generate_block_causal_attention_mask((sequence_length, sequence_length), block_size) + + np.testing.assert_array_equal(splash_mask[:, :], expected) + np.testing.assert_array_equal(dense_mask, expected) + self.assertTrue(expected[1, 3]) + self.assertFalse(expected[3, 4]) + self.assertTrue(expected[4, 3]) + self.assertTrue(expected[8, 9]) + + def test_dot_product_mask_respects_packed_segments(self): + sequence_length = 12 + query = jnp.zeros((1, sequence_length, 1, 8)) + key = jnp.zeros((1, sequence_length, 1, 8)) + segment_ids = jnp.asarray([[1] * 8 + [2] * 4], dtype=jnp.int32) + + mask = self._make_op(sequence_length).generate_attention_mask( + query, + key, + segment_ids, + MODEL_MODE_TRAIN, + ) + + positions = np.arange(sequence_length) + expected = (positions[:, None] // 4 >= positions[None, :] // 4) & ( + np.asarray(segment_ids[0])[:, None] == np.asarray(segment_ids[0])[None, :] + ) + np.testing.assert_array_equal(np.asarray(mask == 0.0)[0, 0, 0], expected) + + def test_autoregressive_mask_is_unchanged(self): + key_length = 8 + query = jnp.zeros((1, 1, 1, 8)) + key = jnp.zeros((1, key_length, 1, 8)) + segment_ids = jnp.asarray([[1, 1, 0, 1, 0, 0, 1, 0]], dtype=jnp.int32) + + default_mask = self._make_op(key_length, attention_type=None).generate_attention_mask( + query, + key, + segment_ids, + MODEL_MODE_AUTOREGRESSIVE, + ) + block_diffusion_mask = self._make_op(key_length).generate_attention_mask( + query, + key, + segment_ids, + MODEL_MODE_AUTOREGRESSIVE, + ) + + np.testing.assert_array_equal(block_diffusion_mask, default_mask) + np.testing.assert_array_equal( + np.asarray(default_mask == 0.0)[0, 0, 0, 0], + np.asarray(segment_ids[0] == DECODING_ACTIVE_SEQUENCE_INDICATOR), + ) + + def test_invalid_block_size_raises(self): + with self.assertRaises(ValueError): + BlockCausalMask((4, 4), 0) + with self.assertRaises(ValueError): + _generate_block_causal_attention_mask((4, 4), -1) + + +class AttentionTypeResolutionTest(unittest.TestCase): + + def test_config_selects_block_diffusion_without_model_dispatch(self): + config = types.SimpleNamespace(attention_type=AttentionType.BLOCK_DIFFUSION.value) + + self.assertEqual(_resolve_attention_type(config, None), AttentionType.BLOCK_DIFFUSION) + self.assertEqual(_resolve_attention_type(config, AttentionType.GLOBAL), AttentionType.BLOCK_DIFFUSION) + self.assertEqual(_resolve_attention_type(config, AttentionType.FULL), AttentionType.FULL) + self.assertEqual(_resolve_attention_type(config, AttentionType.LOCAL_SLIDING), AttentionType.LOCAL_SLIDING) + self.assertEqual( + _resolve_attention_type(types.SimpleNamespace(attention_type=AttentionType.GLOBAL.value), AttentionType.FULL), + AttentionType.FULL, + ) + self.assertEqual(_resolve_attention_type(types.SimpleNamespace(), None), AttentionType.GLOBAL) + + def test_attention_op_honors_config_without_overriding_specialized_layers(self): + config = types.SimpleNamespace( + attention_type=AttentionType.BLOCK_DIFFUSION.value, + block_diffusion_block_size=4, + ) + op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=8, + mesh=types.SimpleNamespace(shape={}), + attention_kernel="dot_product", + attention_type=AttentionType.GLOBAL, + ) + + self.assertEqual(op.attention_type, AttentionType.BLOCK_DIFFUSION) + + full_op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=8, + mesh=types.SimpleNamespace(shape={}), + attention_kernel="dot_product", + attention_type=AttentionType.FULL, + ) + + self.assertEqual(full_op.attention_type, AttentionType.FULL) + + class LoadBalancedMaskTest(unittest.TestCase): """Tests for load-balanced Splash masks.""" @@ -370,6 +510,20 @@ def test_load_balanced_chunk_window(self): np.testing.assert_array_equal((causal_mask & chunk_mask)[:, :], expected_mask) + def test_load_balanced_block_causal_mask(self): + sequence_length = 8 + block_size = 2 + mask = attention_op.LoadBalancedBlockCausalMask( + shape=(sequence_length, sequence_length), + block_size=block_size, + cp_size=2, + ) + query_positions = mask.q_sequence[:, None] + key_positions = np.arange(sequence_length)[None, :] + expected = query_positions // block_size >= key_positions // block_size + + np.testing.assert_array_equal(mask[:, :], expected) + def test_dot_product_local_mask_uses_segment_positions(self): config = types.SimpleNamespace(context_parallel_load_balance=True, context_sharding="context") mesh = types.SimpleNamespace(shape={"context": 4}) diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index f0504e4dd8..f333748192 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -192,6 +192,54 @@ def test_load_balanced_chunk_context_parallel_config(self): self.assertEqual(config.attention_type, "chunk") self.assertTrue(config.context_parallel_load_balance) + def test_block_diffusion_config(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "attention=dot_product", + "attention_type=block_diffusion", + "block_diffusion_block_size=7", + "vocab_size=256", + "max_target_length=2048", + "hardware=cpu", + ] + config = pyconfig.initialize(argv) + + self.assertEqual(config.attention_type, "block_diffusion") + self.assertEqual(config.block_diffusion_block_size, 7) + self.assertNotEqual(config.max_target_length % config.block_diffusion_block_size, 0) + self.assertFalse(hasattr(config, "enable_block_diffusion")) + + def test_block_diffusion_invalid_configs_raise(self): + base_overrides = { + "run_name": "test", + "steps": 1, + "attention": "dot_product", + "attention_type": "block_diffusion", + "block_diffusion_block_size": 32, + "vocab_size": 256, + "max_target_length": 2048, + "hardware": "cpu", + } + cases = [ + {"block_diffusion_block_size": 0}, + {"attention": "cudnn_flash_te"}, + {"attention": "flash", "hardware": "gpu"}, + ] + for overrides in cases: + with self.subTest(overrides=overrides): + config = base_overrides | overrides + argv = ["", _BASE_CONFIG_PATH, *(f"{key}={value}" for key, value in config.items())] + with self.assertRaises((ValueError, pydantic.ValidationError)): + pyconfig.initialize(argv) + + def test_default_attention_remains_global(self): + config = pyconfig.initialize(["", _BASE_CONFIG_PATH, "run_name=test", "steps=1"]) + + self.assertEqual(config.attention_type, "global") + @unittest.mock.patch.dict(os.environ, {pyconfig.yaml_key_to_env_key("steps"): "123"}) def test_env_override(self): """Tests that environment variables override YAML values.""" diff --git a/tests/unit/tokamax_splash_attention_mask_test.py b/tests/unit/tokamax_splash_attention_mask_test.py index 590e8046b8..b1fe154029 100644 --- a/tests/unit/tokamax_splash_attention_mask_test.py +++ b/tests/unit/tokamax_splash_attention_mask_test.py @@ -20,6 +20,7 @@ from maxtext.kernels.tokamax_splash_attention import splash_attention_mask from maxtext.kernels.tokamax_splash_attention import splash_attention_mask_info +from maxtext.layers.attention_op import BlockCausalMask class TokamaxSplashAttentionMaskTest(absltest.TestCase): @@ -56,6 +57,22 @@ def test_process_causal_mask_keeps_lazy_mask_function(self): self.assertIsNone(mask_info.partial_mask_blocks) np.testing.assert_array_equal(mask_info.q_sequence, np.arange(8, dtype=np.int32)) + def test_process_block_causal_mask_keeps_lazy_mask_function(self): + mask = BlockCausalMask((8, 8), block_size=4) + + mask_info, mask_function = splash_attention_mask_info.process_mask( + mask, + block_shape=(2, 2), + q_seq_shards=1, + kv_seq_shards=1, + return_dynamic_grid=True, + ) + + self.assertIs(mask_function, mask.mask_function) + self.assertIsNotNone(mask_info.q_sequence) + self.assertIsNone(mask_info.partial_mask_blocks) + np.testing.assert_array_equal(mask_info.q_sequence, np.arange(8, dtype=np.int32)) + if __name__ == "__main__": absltest.main() From 1822108c6dd30d438cca3eb7b551b4c1693add5d Mon Sep 17 00:00:00 2001 From: ethannnnnn Date: Wed, 22 Jul 2026 23:54:36 -0700 Subject: [PATCH 2/6] [maxtext] Add generalized block-diffusion CFT and SFT Add a model-agnostic block-diffusion training objective without changing the causal LM default. The data pipeline emits separate validity, completion, corruption, and loss masks, supports same-position/all-masked and shifted/seeded canvases, and preserves those fields through context-parallel reordering and shaped batches. Align model logits to physical targets in a shared scoring utility and consume explicit loss weights in Linen, NNX, native gradient accumulation, and the Tunix SFT adapter. The adapter composes the stock Tunix PeftTrainer and its LossOutput contract, including denominator-aware accumulation and evaluation; the existing MaxText AR trainer path remains unchanged. Tests cover partial blocks, prompt protection, deterministic corruption, alignment after sequence reordering, strict mask requirements, zero-weight gradients, causal no-regression, weighted accumulation/evaluation acceptance, and SFT adapter validation. Test Plan: - 196 passed, 14 skipped, 3 deselected; 66 subtests passed in focused MaxText unit suite - 11 passed in Tunix-backed post-training SFT suite - Pylint 10.00/10 on new scoring and loss/GA tests - Pyink, yamllint, codespell, pycompile, and git diff checks pass --- src/maxtext/configs/base.yml | 5 + src/maxtext/configs/types.py | 52 ++++ src/maxtext/diffusion/__init__.py | 15 + src/maxtext/diffusion/scoring.py | 43 +++ .../input_pipeline/hf_data_processing.py | 51 +++- .../input_pipeline/input_pipeline_utils.py | 117 +++++++- .../synthetic_data_processing.py | 5 + .../integration/tunix/diffusion_sft.py | 180 ++++++++++++ src/maxtext/trainers/post_train/sft/hooks.py | 2 +- .../trainers/post_train/sft/train_sft.py | 34 ++- src/maxtext/trainers/pre_train/train.py | 51 +++- src/maxtext/utils/gradient_accumulation.py | 10 +- src/maxtext/utils/max_utils.py | 3 + src/maxtext/utils/maxtext_utils.py | 4 + .../post_training/unit/diffusion_sft_test.py | 95 +++++++ tests/post_training/unit/sft_hooks_test.py | 12 + tests/post_training/unit/train_sft_test.py | 114 ++++++++ tests/unit/configs_value_test.py | 101 +++++++ tests/unit/gradient_accumulation_nnx_test.py | 32 +++ tests/unit/hf_data_processing_test.py | 72 +++++ tests/unit/input_pipeline_utils_test.py | 214 +++++++++++++- tests/unit/max_utils_test.py | 27 ++ tests/unit/maxtext_utils_test.py | 28 +- tests/unit/pre_train_loss_mask_test.py | 264 ++++++++++++++++++ tests/unit/synthetic_data_test.py | 23 +- 25 files changed, 1530 insertions(+), 24 deletions(-) create mode 100644 src/maxtext/diffusion/__init__.py create mode 100644 src/maxtext/diffusion/scoring.py create mode 100644 src/maxtext/integration/tunix/diffusion_sft.py create mode 100644 tests/post_training/unit/diffusion_sft_test.py create mode 100644 tests/unit/pre_train_loss_mask_test.py diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 4d0bf560e8..1091ff7d8d 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -791,6 +791,11 @@ olmo_apply_ngram_filter: true # mask instances with repetitive n-grams (OLMo-cor # Training loop steps: 150_001 # If set to -1 then will inherit value from learning_rate_schedule_steps log_period: 100 # The frequency of Tensorboard flush, gcs metrics writing, and managed profiler metrics updating. +training_objective: 'causal_lm' # Supported objectives: causal_lm, block_diffusion +block_diffusion_mask_id: -1 # Tokenizer mask-token id; required for training_objective='block_diffusion' +block_diffusion_min_noise: 0.001 # Minimum per-block corruption probability for block-diffusion training +block_diffusion_logit_alignment: 'same_position' # Supported alignments: same_position, shifted +block_diffusion_canvas_policy: 'all_masked' # Supported canvases: all_masked, seed_and_mask jax_distributed_initialization_timeout: 300 # This is the default timeout in https://github.com/jax-ml/jax/blob/main/jax/_src/distributed.py # Note there are two separate initializations - the jax coordination service (aka jax.distributed.initialize) and the backend (e.g. PjRT), the timeout above refers diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 8e42109894..ccc817b090 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1583,6 +1583,28 @@ class Distillation(BaseModel): class TrainingLoop(BaseModel): """Configuration for the main training loop, evaluation, and reproducibility.""" + training_objective: Literal["causal_lm", "block_diffusion"] = Field( + "causal_lm", + description="The token-prediction objective used to prepare targets and compute loss.", + ) + block_diffusion_mask_id: int = Field( + -1, + description="The tokenizer mask-token id required by the block-diffusion training objective.", + ) + block_diffusion_min_noise: float = Field( + 1.0e-3, + gt=0.0, + le=1.0, + description="The minimum corruption probability sampled independently for each block.", + ) + block_diffusion_logit_alignment: Literal["same_position", "shifted"] = Field( + "same_position", + description="How model logits align to clean target-token positions.", + ) + block_diffusion_canvas_policy: Literal["all_masked", "seed_and_mask"] = Field( + "all_masked", + description="Whether every block is fully maskable or begins with a clean anchor token.", + ) steps: int = Field( 150_001, ge=-1, @@ -3295,6 +3317,36 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "Block-diffusion attention with attention='autoselected' or attention='flash' requires hardware='tpu'; " "use attention='dot_product' on other hardware." ) + if self.training_objective == "block_diffusion": + if self.attention_type != AttentionType.BLOCK_DIFFUSION.value: + raise ValueError("`training_objective='block_diffusion'` requires `attention_type='block_diffusion'`.") + if self.block_diffusion_mask_id < 0 or self.block_diffusion_mask_id >= self.vocab_size: + raise ValueError( + f"`block_diffusion_mask_id` ({self.block_diffusion_mask_id}) must satisfy " + f"0 <= block_diffusion_mask_id < vocab_size ({self.vocab_size})." + ) + if self.packing: + raise ValueError("`training_objective='block_diffusion'` requires `packing=False`.") + if self.mtp_num_layers > 0: + raise ValueError("`training_objective='block_diffusion'` is not compatible with MTP.") + if self.num_vocab_tiling > 1: + raise ValueError("`training_objective='block_diffusion'` is not compatible with vocabulary tiling.") + if self.dataset_type != "hf": + raise ValueError("`training_objective='block_diffusion'` currently requires `dataset_type='hf'`.") + if self.use_dpo: + raise ValueError("`training_objective='block_diffusion'` is not compatible with DPO.") + if self.use_multimodal or self.use_audio: + raise ValueError("`training_objective='block_diffusion'` currently supports text-only training.") + valid_model_contracts = { + ("same_position", "all_masked"), + ("shifted", "seed_and_mask"), + } + model_contract = (self.block_diffusion_logit_alignment, self.block_diffusion_canvas_policy) + if model_contract not in valid_model_contracts: + raise ValueError( + "Block-diffusion training supports only `same_position/all_masked` or `shifted/seed_and_mask`; " + f"received `{model_contract[0]}/{model_contract[1]}`." + ) if self.quantize_kvcache and not self.kv_quant_axis: raise ValueError("`kv_quant_axis` cannot be empty when quantize_kvcache is True.") if ( diff --git a/src/maxtext/diffusion/__init__.py b/src/maxtext/diffusion/__init__.py new file mode 100644 index 0000000000..2c9fc63d21 --- /dev/null +++ b/src/maxtext/diffusion/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared block-diffusion model semantics.""" diff --git a/src/maxtext/diffusion/scoring.py b/src/maxtext/diffusion/scoring.py new file mode 100644 index 0000000000..bf220219dc --- /dev/null +++ b/src/maxtext/diffusion/scoring.py @@ -0,0 +1,43 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Target-alignment utilities shared by diffusion training objectives.""" + +import jax.numpy as jnp + + +def align_logits_to_targets(logits, alignment, positions=None, validity_mask=None): + """Aligns model output positions with clean block-diffusion targets.""" + if alignment == "same_position": + return logits + if alignment == "shifted": + if positions is None: + indices = jnp.maximum(jnp.arange(logits.shape[1], dtype=jnp.int32) - 1, 0) + return logits[:, indices, :] + positions = jnp.asarray(positions, dtype=jnp.int32) + if positions.shape != logits.shape[:2]: + raise ValueError(f"positions must match logits [batch, length]; got {positions.shape} and {logits.shape[:2]}") + if validity_mask is None: + validity_mask = jnp.ones_like(positions, dtype=jnp.bool_) + sequence_length = logits.shape[1] + array_indices = jnp.broadcast_to(jnp.arange(sequence_length, dtype=jnp.int32), positions.shape) + row_indices = jnp.broadcast_to(jnp.arange(logits.shape[0], dtype=jnp.int32)[:, None], positions.shape) + safe_positions = jnp.clip(positions, 0, sequence_length - 1) + updates = jnp.where(validity_mask, array_indices + 1, 0) + position_to_index = jnp.zeros_like(positions).at[row_indices, safe_positions].max(updates) + previous_positions = jnp.maximum(positions - 1, 0) + source_indices = jnp.take_along_axis(position_to_index, jnp.clip(previous_positions, 0, sequence_length - 1), axis=1) + source_indices = jnp.maximum(source_indices - 1, 0) + return jnp.take_along_axis(logits, source_indices[..., None], axis=1) + raise ValueError(f"Unsupported block-diffusion logit alignment: {alignment}") diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py index 370f1895bd..a047baadbb 100644 --- a/src/maxtext/input_pipeline/hf_data_processing.py +++ b/src/maxtext/input_pipeline/hf_data_processing.py @@ -42,6 +42,40 @@ def _get_pad_id(tokenizer): return pad_id +def _get_training_objective_transform( + config, + *, + shift, + use_dpo, + use_sft, + completion_only, + packing, + pad_id, + bos_token_id, +): + """Selects the aligned or next-token target transform for an HF pipeline.""" + objective = getattr(config, "training_objective", "causal_lm") + if objective == "block_diffusion": + if packing: + raise ValueError("Block-diffusion training requires packing=False.") + if use_dpo: + raise ValueError("Block-diffusion training is not compatible with DPO.") + return input_pipeline_utils.BlockDiffusionCorruption( + block_size=config.block_diffusion_block_size, + mask_id=config.block_diffusion_mask_id, + min_noise=config.block_diffusion_min_noise, + axis=1, + completion_only=bool(use_sft and completion_only), + seed_first_token=config.block_diffusion_canvas_policy == "seed_and_mask", + include_seed_in_loss=config.block_diffusion_logit_alignment == "shifted", + ) + if objective != "causal_lm": + raise ValueError(f"Unsupported training objective: {objective}") + if shift and not use_dpo: + return input_pipeline_utils.ShiftData(ignored_ids=[pad_id, bos_token_id], axis=1) + return None + + def vision_sft_preprocessing_pipeline( dataset, config, @@ -321,12 +355,15 @@ def preprocessing_pipeline( ) operations = [] if use_sft: + is_block_diffusion = getattr(config, "training_objective", "causal_lm") == "block_diffusion" operations.append( input_pipeline_utils.SFTPromptMasking( text_column_name=data_column_names[0], completion_only=sft_train_on_completion_only, max_target_length=max_target_length, unk_id=pad_id, + target_aligned=is_block_diffusion, + emit_completion_mask=is_block_diffusion, ) ) data_column_names = ("inputs", "targets") @@ -357,8 +394,18 @@ def preprocessing_pipeline( operations.append(input_pipeline_utils.PadOrTrimToMaxLength(max_target_length, pad_id)) operations.append(grain.Batch(batch_size=batch_size, drop_remainder=drop_remainder)) - if shift and not use_dpo: - operations.append(input_pipeline_utils.ShiftData(ignored_ids=[pad_id, tokenizer.bos_token_id], axis=1)) + target_transform = _get_training_objective_transform( + config, + shift=shift, + use_dpo=use_dpo, + use_sft=use_sft, + completion_only=sft_train_on_completion_only, + packing=packing, + pad_id=pad_id, + bos_token_id=tokenizer.bos_token_id, + ) + if target_transform is not None: + operations.append(target_transform) # Since HuggingFace IterableDataset does not support access through index # Indexes generated by dummy_index_sampler is not used. diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index d996a88908..6b294a6dcb 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -325,14 +325,25 @@ def tokenization(example, hf_tokenizer, truncation, max_length, column_names): @dataclasses.dataclass class SFTPromptMasking(grain.MapTransform): """Construct inputs and targets for SFT training. Concat prompt and completion to generate inputs. - For targets, if train on completion only, the prompt will be masked by unk_id. Otherwise the same as inputs. + Causal SFT can mask prompt targets with ``unk_id``. Target-aligned objectives + keep clean targets and may emit role-derived completion metadata instead. """ - def __init__(self, text_column_name, completion_only, max_target_length, unk_id=0): + def __init__( + self, + text_column_name, + completion_only, + max_target_length, + unk_id=0, + target_aligned=False, + emit_completion_mask=False, + ): self.text_column_name = text_column_name self.completion_only = completion_only self.max_target_length = max_target_length self.unk_id = unk_id + self.target_aligned = target_aligned + self.emit_completion_mask = emit_completion_mask def map(self, element): """ @@ -340,18 +351,24 @@ def map(self, element): It concatenates the prompt and completion to form the `inputs` sequence. For the `targets` sequence: - If `self.completion_only` is `True`, the prompt portion of the - concatenated sequence is masked using `self.unk_id`. + concatenated sequence is masked using `self.unk_id`, unless + `self.target_aligned` requests clean same-position targets. - If `self.completion_only` is `False`, the target sequence is identical to the input sequence. """ - inputs, targets = [], [] + inputs, targets, completion_mask = [], [], [] for i, text in enumerate(element[self.text_column_name]): inputs += text - targets += [self.unk_id] * len(text) if self.completion_only and element["is_prompt"][i] else text - return { + is_prompt = element["is_prompt"][i] + targets += [self.unk_id] * len(text) if self.completion_only and is_prompt and not self.target_aligned else text + completion_mask += [not is_prompt] * len(text) + output = { "inputs": np.asarray(inputs[: self.max_target_length], dtype=np.int32), "targets": np.asarray(targets[: self.max_target_length], dtype=np.int32), } + if self.emit_completion_mask: + output["completion_mask"] = np.asarray(completion_mask[: self.max_target_length], dtype=np.int32) + return output @dataclasses.dataclass @@ -838,7 +855,7 @@ def map( """map to each element""" data_columns = list(element.keys()) for data_column in data_columns: - if data_column != "images": + if data_column not in ("images", "completion_mask"): if isinstance(element[data_column], mm_utils.PreprocessorOutput): raise TypeError("Only 'images' column can be of type PreprocessorOutput.") @@ -863,6 +880,8 @@ def map( element["images"] = self._pad_image_and_mask(element["images"]) # pyrefly: ignore[bad-argument-type] + elif key == "completion_mask": + element[key] = self._pad_text(element[key], self.max_length, 0) # pyrefly: ignore[bad-argument-type] elif "true_length" not in key: element[key] = self._pad_text(element[key], self.max_length, self.pad_id) # pyrefly: ignore[bad-argument-type] return element @@ -986,6 +1005,90 @@ def map(self, element): return shift_and_refine(element, ignored_ids=self.ignored_ids, axis=self.axis) +@dataclasses.dataclass +class BlockDiffusionCorruption(grain.RandomMapTransform): + """Builds explicit completion, corruption, and loss masks for block diffusion.""" + + def __init__( + self, + block_size: int, + mask_id: int, + min_noise: float = 1.0e-3, + axis: int = 1, + completion_only: bool = False, + seed_first_token: bool = False, + include_seed_in_loss: bool = False, + ): + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + if not 0.0 < min_noise <= 1.0: + raise ValueError(f"min_noise must satisfy 0 < min_noise <= 1, got {min_noise}") + if include_seed_in_loss and not seed_first_token: + raise ValueError("include_seed_in_loss requires seed_first_token=True") + self.block_size = block_size + self.mask_id = mask_id + self.min_noise = min_noise + self.axis = axis + self.completion_only = completion_only + self.seed_first_token = seed_first_token + self.include_seed_in_loss = include_seed_in_loss + + def random_map(self, element, rng: np.random.Generator): + inputs = np.asarray(element["inputs"]) + targets = np.asarray(element["targets"]) + targets_segmentation = np.asarray(element["targets_segmentation"]) + if inputs.shape != targets.shape or inputs.shape != targets_segmentation.shape: + raise ValueError( + "inputs, targets, and targets_segmentation must have identical shapes, got " + f"{inputs.shape}, {targets.shape}, and {targets_segmentation.shape}" + ) + if inputs.ndim == 0 or not -inputs.ndim <= self.axis < inputs.ndim: + raise ValueError(f"axis {self.axis} is invalid for inputs with {inputs.ndim} dimensions") + + axis = self.axis % inputs.ndim + validity_mask = targets_segmentation != 0 + if self.completion_only and "completion_mask" not in element: + raise ValueError("completion_only block diffusion requires an explicit completion_mask") + completion_mask = np.asarray(element.get("completion_mask", validity_mask)) != 0 + if completion_mask.shape != inputs.shape: + raise ValueError(f"completion_mask must match inputs shape; received {completion_mask.shape} and {inputs.shape}") + completion_mask &= validity_mask + eligible_mask = completion_mask if self.completion_only else validity_mask + + moved_eligible = np.moveaxis(eligible_mask, axis, -1) + rows = moved_eligible.reshape(-1, inputs.shape[axis]) + block_count = (rows.shape[-1] + self.block_size - 1) // self.block_size + padded_length = block_count * self.block_size + padded_rows = np.pad(rows, ((0, 0), (0, padded_length - rows.shape[-1]))) + eligible_blocks = padded_rows.reshape(rows.shape[0], block_count, self.block_size) + seed_loss_blocks = np.zeros_like(eligible_blocks) + if self.seed_first_token: + if self.include_seed_in_loss: + seed_loss_blocks = eligible_blocks & (np.arange(self.block_size) == 0) + seed_loss_blocks &= np.arange(block_count)[None, :, None] != 0 + eligible_blocks = eligible_blocks & (np.arange(self.block_size) != 0) + + noise = rng.uniform(self.min_noise, 1.0, size=eligible_blocks.shape[:-1] + (1,)) + selected_blocks = eligible_blocks & (rng.random(size=eligible_blocks.shape) < noise) + needs_fallback = eligible_blocks.any(axis=-1) & ~selected_blocks.any(axis=-1) + fallback_scores = np.where(eligible_blocks, rng.random(size=eligible_blocks.shape), -1.0) + fallback_offsets = np.argmax(fallback_scores, axis=-1) + fallback = np.arange(self.block_size) == fallback_offsets[..., None] + selected_blocks |= needs_fallback[..., None] & fallback + + corruption_rows = selected_blocks.reshape(rows.shape[0], padded_length)[:, : rows.shape[-1]] + loss_rows = (selected_blocks | seed_loss_blocks).reshape(rows.shape[0], padded_length)[:, : rows.shape[-1]] + corruption_mask = np.moveaxis(corruption_rows.reshape(moved_eligible.shape), -1, axis) + targets_loss_mask = np.moveaxis(loss_rows.reshape(moved_eligible.shape), -1, axis) + output = dict(element) + output["inputs"] = np.where(corruption_mask, inputs.dtype.type(self.mask_id), inputs) + output["targets"] = targets + output["completion_mask"] = completion_mask.astype(targets_segmentation.dtype) + output["corruption_mask"] = corruption_mask.astype(targets_segmentation.dtype) + output["targets_loss_mask"] = targets_loss_mask.astype(targets_segmentation.dtype) + return output + + @dataclasses.dataclass class ComputeQwen3OmniPositions(grain.MapTransform): """Computes 3D position IDs for Qwen3-Omni multimodal sequences. diff --git a/src/maxtext/input_pipeline/synthetic_data_processing.py b/src/maxtext/input_pipeline/synthetic_data_processing.py index d10816435c..a184e522cd 100644 --- a/src/maxtext/input_pipeline/synthetic_data_processing.py +++ b/src/maxtext/input_pipeline/synthetic_data_processing.py @@ -111,6 +111,11 @@ def get_place_holder_synthetic_data(config: pyconfig.HyperParameters): "targets_position": neg_ones, "targets_segmentation": neg_ones, } + if getattr(config, "training_objective", "causal_lm") == "block_diffusion": + zeros = np.zeros_like(neg_ones) + batch["completion_mask"] = zeros + batch["corruption_mask"] = zeros + batch["targets_loss_mask"] = zeros def infinite_iterator(): while True: diff --git a/src/maxtext/integration/tunix/diffusion_sft.py b/src/maxtext/integration/tunix/diffusion_sft.py new file mode 100644 index 0000000000..b85132db38 --- /dev/null +++ b/src/maxtext/integration/tunix/diffusion_sft.py @@ -0,0 +1,180 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MaxText adapters for target-aligned Tunix diffusion SFT.""" + +from collections.abc import Mapping +from typing import Any + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.diffusion import scoring +from tunix.diffusion import types as diffusion_types + + +_REQUIRED_FIELDS = ( + "inputs", + "inputs_position", + "inputs_segmentation", + "targets", + "targets_position", + "targets_segmentation", + "completion_mask", + "corruption_mask", + "targets_loss_mask", +) + + +def _concrete_numpy(value): + if isinstance(value, jax.core.Tracer): + return None + if isinstance(value, jax.Array) and not value.is_fully_addressable: + return None + return np.asarray(value) + + +def _validate_batch_masks( + positions, + validity_mask, + completion_mask, + corruption_mask, + loss_weights, + *, + alignment, + block_size, +): + shapes = { + "positions": tuple(positions.shape), + "validity_mask": tuple(validity_mask.shape), + "completion_mask": tuple(completion_mask.shape), + "corruption_mask": tuple(corruption_mask.shape), + "targets_loss_mask": tuple(loss_weights.shape), + } + if len(set(shapes.values())) != 1: + raise ValueError( + f"diffusion SFT masks must have identical shapes; received {shapes}" + ) + concrete = [ + _concrete_numpy(value) + for value in ( + positions, + validity_mask, + completion_mask, + corruption_mask, + loss_weights, + ) + ] + if any(value is None for value in concrete): + return + concrete_positions = np.asarray(concrete[0]) + validity, completion, corruption = ( + np.asarray(value, dtype=bool) for value in concrete[1:4] + ) + weights = np.asarray(concrete[4]) + weighted = weights != 0 + if not np.all(np.isfinite(weights)) or np.any(weights < 0): + raise ValueError("targets_loss_mask must contain finite nonnegative weights") + if np.any(completion & ~validity): + raise ValueError("completion_mask must be a subset of valid target positions") + if np.any(corruption & ~completion): + raise ValueError("corruption_mask must be a subset of completion_mask") + if np.any(weighted & ~completion): + raise ValueError("targets_loss_mask must be a subset of completion_mask") + allowed = corruption.copy() + if alignment == "shifted": + allowed |= (concrete_positions > 0) & (concrete_positions % block_size == 0) + if np.any(weighted & ~allowed): + raise ValueError( + "diffusion SFT loss weights must own corrupted targets or shifted block anchors" + ) + + +def create_batch_adapter(config): + """Builds a raw MaxText batch adapter for Tunix diffusion SFT.""" + alignment = config.block_diffusion_logit_alignment + block_size = int(config.block_diffusion_block_size) + + def adapt(raw_batch: Mapping[str, Any]) -> diffusion_types.DiffusionTokenBatch: + missing = sorted(name for name in _REQUIRED_FIELDS if name not in raw_batch) + if missing: + raise ValueError( + f"block-diffusion SFT requires explicit batch fields; missing {missing}" + ) + targets = jnp.asarray(raw_batch["targets"]) + positions = jnp.asarray(raw_batch["targets_position"], dtype=jnp.int32) + validity_mask = jnp.asarray(raw_batch["targets_segmentation"]) != 0 + completion_mask = jnp.asarray(raw_batch["completion_mask"], dtype=jnp.bool_) + corruption_mask = jnp.asarray(raw_batch["corruption_mask"], dtype=jnp.bool_) + raw_loss_weights = jnp.asarray( + raw_batch["targets_loss_mask"], dtype=jnp.float32 + ) + _validate_batch_masks( + positions, + validity_mask, + completion_mask, + corruption_mask, + raw_loss_weights, + alignment=alignment, + block_size=block_size, + ) + allowed = corruption_mask + if alignment == "shifted": + allowed |= (positions > 0) & (positions % block_size == 0) + loss_weights = jnp.where( + validity_mask & completion_mask & allowed, raw_loss_weights, 0.0 + ) + return diffusion_types.DiffusionTokenBatch.create( + model_inputs={ + "input_tokens": jnp.asarray(raw_batch["inputs"]), + "input_positions": jnp.asarray( + raw_batch["inputs_position"], dtype=jnp.int32 + ), + "input_segmentation": jnp.asarray(raw_batch["inputs_segmentation"]), + "targets": targets, + "target_positions": positions, + "target_segmentation": jnp.asarray(raw_batch["targets_segmentation"]), + }, + target_ids=targets, + loss_weights=loss_weights, + ) + + return adapt + + +def create_target_aligned_logits_fn(config): + """Builds a MaxText scorer satisfying Tunix's diffusion logits contract.""" + alignment = config.block_diffusion_logit_alignment + enable_dropout = bool(config.enable_dropout) + + def logits_fn(model: nnx.Module, model_inputs: diffusion_types.ModelInputs): + base_model = getattr(model, "base", model) + logits = base_model( + decoder_input_tokens=model_inputs["input_tokens"], + decoder_positions=model_inputs["input_positions"], + decoder_segment_ids=model_inputs["input_segmentation"], + enable_dropout=enable_dropout, + decoder_target_tokens=model_inputs["targets"], + decoder_target_mask=model_inputs["target_segmentation"], + ) + return scoring.align_logits_to_targets( + logits, + alignment, + model_inputs["target_positions"], + model_inputs["target_segmentation"] != 0, + ) + + return logits_fn diff --git a/src/maxtext/trainers/post_train/sft/hooks.py b/src/maxtext/trainers/post_train/sft/hooks.py index 9f30a3d6d8..03edfbcdb9 100644 --- a/src/maxtext/trainers/post_train/sft/hooks.py +++ b/src/maxtext/trainers/post_train/sft/hooks.py @@ -29,7 +29,7 @@ class SFTTrainingHooks(BaseTrainingHooks): @override def get_total_weights(self, batch) -> jax.Array: """Calculate the number of non-padded tokens in the batch.""" - return jnp.sum(batch["targets_segmentation"] != 0) + return jnp.sum(batch.get("targets_loss_mask", batch["targets_segmentation"]) != 0) class SFTDataHooks(BaseDataHooks): diff --git a/src/maxtext/trainers/post_train/sft/train_sft.py b/src/maxtext/trainers/post_train/sft/train_sft.py index 9002caa16c..715585ec90 100644 --- a/src/maxtext/trainers/post_train/sft/train_sft.py +++ b/src/maxtext/trainers/post_train/sft/train_sft.py @@ -49,10 +49,12 @@ from orbax import checkpoint as ocp +from tunix.sft import diffusion as tunix_diffusion_sft from tunix.sft import metrics_logger, peft_trainer, profiler from maxtext.optimizers import optimizers from maxtext.configs import pyconfig +from maxtext.integration.tunix import diffusion_sft as maxtext_diffusion_sft from maxtext.trainers.pre_train.train import loss_fn from maxtext.common.goodput import ( GoodputEvent, @@ -210,6 +212,9 @@ def loss_func( targets, targets_position, targets_segmentation, + completion_mask=None, + corruption_mask=None, + targets_loss_mask=None, ): data = { "inputs": inputs, @@ -219,12 +224,36 @@ def loss_func( "targets_position": targets_position, "targets_segmentation": targets_segmentation, } + if completion_mask is not None: + data["completion_mask"] = completion_mask + if corruption_mask is not None: + data["corruption_mask"] = corruption_mask + if targets_loss_mask is not None: + data["targets_loss_mask"] = targets_loss_mask return loss_fn(model, mt_config, data, dropout_rng=None, params=None, is_train=True) trainer = trainer.with_loss_fn(loss_func, has_aux=True) return trainer +def configure_training_objective(trainer, mt_config): + """Configures AR or target-aligned diffusion loss without changing AR defaults.""" + if getattr(mt_config, "training_objective", "causal_lm") != "block_diffusion": + return use_maxtext_loss_function(trainer, mt_config) + max_logging.log("Configuring Tunix target-aligned block-diffusion SFT adapter.") + return tunix_diffusion_sft.configure_diffusion_sft( + trainer, + maxtext_diffusion_sft.create_batch_adapter(mt_config), + maxtext_diffusion_sft.create_target_aligned_logits_fn(mt_config), + ) + + +def _create_trainer(model, optimizer, tunix_config, mt_config): + if getattr(mt_config, "training_objective", "causal_lm") == "block_diffusion": + return peft_trainer.PeftTrainer(model, optimizer, tunix_config) + return MaxTextPeftTrainer(model, optimizer, tunix_config) + + def validate_config(config): """Validates the configuration parameters for SFT training.""" if config.optimizer_memory_host_offload: @@ -237,6 +266,7 @@ def validate_config(config): def setup_trainer_state(mt_config, goodput_recorder=None): """Set up prerequisites for training loop.""" + validate_config(mt_config) tunix_config = get_tunix_config(mt_config) with maybe_record_goodput(goodput_recorder, GoodputEvent.TPU_INIT): @@ -261,12 +291,12 @@ def setup_trainer_state(mt_config, goodput_recorder=None): # Provide rules context so 'norm' is translated to mesh axes during maybe_restore with nn_partitioning.axis_rules(mt_config.logical_axis_rules): - trainer = MaxTextPeftTrainer(model, optimizer, tunix_config) + trainer = _create_trainer(model, optimizer, tunix_config, mt_config) if mt_config.lora.lora_restore_path and trainer.train_steps == 0: lora_utils.restore_lora_from_path(trainer.model, mt_config) trainer.with_training_hooks(training_hooks) trainer.with_data_hooks(data_hooks) - trainer = use_maxtext_loss_function(trainer, mt_config) + trainer = configure_training_objective(trainer, mt_config) return trainer, mesh diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index ee92b93435..f27addd189 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -41,6 +41,7 @@ from flax.nnx import variablelib from maxtext.configs import pyconfig +from maxtext.diffusion import scoring as diffusion_scoring from maxtext.utils.globals import EPS from maxtext.utils import elastic_utils # Placeholder: internal @@ -103,6 +104,22 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr loss: average loss aux: a dictionary including intermediate_outputs, xent_sum, and total_weights """ + is_block_diffusion = getattr(config, "training_objective", "causal_lm") == "block_diffusion" + if getattr(config, "attention_type", "global") == "block_diffusion" and not is_block_diffusion: + raise ValueError( + "Block-diffusion attention requires target-aligned block-diffusion losses; " + "causal next-token labels would leak within a bidirectional block." + ) + if is_block_diffusion: + required_masks = {"completion_mask", "corruption_mask", "targets_loss_mask"} + missing_masks = required_masks - data.keys() + if missing_masks: + raise ValueError(f"Block-diffusion loss requires explicit batch masks; missing {sorted(missing_masks)}") + target_shape = data["targets"].shape + for mask_name in required_masks: + if data[mask_name].shape != target_shape: + raise ValueError(f"{mask_name} must match targets shape; got {data[mask_name].shape} and {target_shape}") + # decimate proportion of data when per_device_batch_size<1 if is_train: for k, v in data.items(): @@ -110,6 +127,14 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr else: for k, v in data.items(): data[k] = v[: config.micro_batch_size_to_eval_on, :] + loss_mask_values = data["targets_loss_mask"] if is_block_diffusion else data["targets_segmentation"] + targets_loss_mask = (loss_mask_values != 0) & (data["targets_segmentation"] != 0) + target_positions = data.get("targets_position", data["inputs_position"]) + if is_block_diffusion and config.block_diffusion_canvas_policy == "seed_and_mask": + if config.block_diffusion_logit_alignment == "shifted": + targets_loss_mask &= target_positions != 0 + else: + targets_loss_mask &= target_positions % config.block_diffusion_block_size != 0 mutable_collections = ["intermediates"] if config.mtp_num_layers > 0 and is_train: # The single model.apply call now triggers the entire chain if MTP is enabled: @@ -161,6 +186,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] xent_sum, total_z_loss = vocab_tiling_linen_loss(hidden_states, data, config, model, params, is_train) else: + if is_block_diffusion: + logits = diffusion_scoring.align_logits_to_targets( + logits, + config.block_diffusion_logit_alignment, + target_positions, + data["targets_segmentation"] != 0, + ) one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) @@ -179,9 +211,8 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr debug_sharding=config.debug_sharding, ) - # Mask out paddings at the end of each example. - xent = xent * (data["targets_segmentation"] != 0) - z_loss = z_loss * (data["targets_segmentation"] != 0) + xent = xent * targets_loss_mask + z_loss = z_loss * targets_loss_mask xent_sum = jnp.sum(xent) total_z_loss = jnp.sum(z_loss) @@ -224,6 +255,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] xent_sum, total_z_loss = vocab_tiling_nnx_loss(model, hidden_states, data, config, is_train) else: + if is_block_diffusion: + logits = diffusion_scoring.align_logits_to_targets( + logits, + config.block_diffusion_logit_alignment, + target_positions, + data["targets_segmentation"] != 0, + ) one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) @@ -242,14 +280,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr debug_sharding=config.debug_sharding, ) - # Mask out paddings at the end of each example. - xent = xent * (data["targets_segmentation"] != 0) - z_loss = z_loss * (data["targets_segmentation"] != 0) + xent = xent * targets_loss_mask + z_loss = z_loss * targets_loss_mask xent_sum = jnp.sum(xent) total_z_loss = jnp.sum(z_loss) - total_weights = jnp.sum(data["targets_segmentation"] != 0) + total_weights = jnp.sum(targets_loss_mask) # If gradient accumulation is enabled, we don't need to divide xent_sum # by total_weights and then multiply the computed gradient by total_weights, # since it's equivalent to computing the gradient from xent_sum. diff --git a/src/maxtext/utils/gradient_accumulation.py b/src/maxtext/utils/gradient_accumulation.py index 35fdf65503..eb0786c812 100644 --- a/src/maxtext/utils/gradient_accumulation.py +++ b/src/maxtext/utils/gradient_accumulation.py @@ -162,12 +162,15 @@ def reshape_to_microbatch_accumulations(batch_arr): grad_and_loss, aux = jax.lax.scan( accumulate_gradient, init_grad_and_loss, data, length=config.gradient_accumulation_steps ) + has_weights = grad_and_loss["total_weights"] > 0 + safe_total_weights = jnp.maximum(grad_and_loss["total_weights"], 1) loss = ( - grad_and_loss["loss"] / grad_and_loss["total_weights"] + grad_and_loss["loss"] / safe_total_weights + grad_and_loss["moe_lb_loss"] / config.gradient_accumulation_steps + grad_and_loss["indexer_loss"] / config.gradient_accumulation_steps + grad_and_loss["mtp_loss"] / config.gradient_accumulation_steps ) + loss = jnp.where(has_weights, loss, 0.0) raw_grads = grad_and_loss["grad"] if data_parallel_active: # Mark the gradients unreduced over the "data" axis now that we're outside the @@ -176,7 +179,10 @@ def reshape_to_microbatch_accumulations(batch_arr): unreduced_shardings = jax.tree.map(update_sharding_for_unreduced, params_shardings) raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, unreduced_shardings) raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, params_shardings) - raw_grads = jax.tree_util.tree_map(lambda arr: arr / grad_and_loss["total_weights"], raw_grads) + raw_grads = jax.tree_util.tree_map( + lambda arr: jnp.where(has_weights, arr / safe_total_weights, jnp.zeros_like(arr)), + raw_grads, + ) aux = jax.tree.map(lambda x: jnp.sum(x, axis=0), aux) # pytype: disable=module-attr if is_nnx: diff --git a/src/maxtext/utils/max_utils.py b/src/maxtext/utils/max_utils.py index e6aebce059..943b1c63c2 100644 --- a/src/maxtext/utils/max_utils.py +++ b/src/maxtext/utils/max_utils.py @@ -957,6 +957,9 @@ def reorder_causal_load_balanced(batch, cp_size, reorder_strategy, hardware="tpu "targets_position", "inputs_segmentation", "targets_segmentation", + "completion_mask", + "corruption_mask", + "targets_loss_mask", } if hardware in ("gpu", "gpu_multiprocess"): diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 6c71e0487e..7f5845a72d 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -161,6 +161,10 @@ def get_shaped_batch(config, batch_sharding=None): shaped_batch["targets"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) shaped_batch["targets_position"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) shaped_batch["targets_segmentation"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) + if getattr(config, "training_objective", "causal_lm") == "block_diffusion": + shaped_batch["completion_mask"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) + shaped_batch["corruption_mask"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) + shaped_batch["targets_loss_mask"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) if config.use_multimodal: image_shape = mm_processor.get_dummy_image_shape_for_init( config.model_name, batch_size=config.micro_batch_size_to_train_on diff --git a/tests/post_training/unit/diffusion_sft_test.py b/tests/post_training/unit/diffusion_sft_test.py new file mode 100644 index 0000000000..475b0ecdd7 --- /dev/null +++ b/tests/post_training/unit/diffusion_sft_test.py @@ -0,0 +1,95 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the MaxText-to-Tunix diffusion SFT adapter.""" + +from types import SimpleNamespace + +import jax.numpy as jnp +import numpy as np +import pytest + +from maxtext.diffusion import scoring +from maxtext.integration.tunix import diffusion_sft + + +pytestmark = [pytest.mark.post_training, pytest.mark.cpu_only] + + +def _config(alignment="shifted"): + return SimpleNamespace( + block_diffusion_block_size=4, + block_diffusion_logit_alignment=alignment, + enable_dropout=False, + ) + + +def _raw_batch(): + positions = jnp.arange(8, dtype=jnp.int32)[None, :] + segmentation = jnp.ones((1, 8), dtype=jnp.int32) + return { + "inputs": jnp.asarray([[10, 11, 99, 99, 14, 99, 99, 99]], dtype=jnp.int32), + "inputs_position": positions, + "inputs_segmentation": segmentation, + "targets": jnp.asarray([[10, 11, 12, 13, 14, 15, 16, 17]], dtype=jnp.int32), + "targets_position": positions, + "targets_segmentation": segmentation, + "completion_mask": jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 1]], dtype=jnp.int32), + "corruption_mask": jnp.asarray([[0, 0, 1, 1, 0, 1, 1, 1]], dtype=jnp.int32), + "targets_loss_mask": jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 1]], dtype=jnp.int32), + } + + +def test_batch_adapter_preserves_shifted_anchor_weight(): + batch = diffusion_sft.create_batch_adapter(_config())(_raw_batch()) + + np.testing.assert_array_equal(batch.target_ids, _raw_batch()["targets"]) + np.testing.assert_array_equal(batch.loss_weights, [[0, 0, 1, 1, 1, 1, 1, 1]]) + assert batch.loss_weights.dtype == jnp.float32 + + +def test_batch_adapter_rejects_unowned_clean_target(): + raw = _raw_batch() + raw["corruption_mask"] = raw["corruption_mask"].at[0, 3].set(0) + + with pytest.raises(ValueError, match="corrupted targets or shifted block anchors"): + diffusion_sft.create_batch_adapter(_config())(raw) + + +def test_logits_adapter_uses_target_alignment(): + raw_logits = jnp.arange(1 * 8 * 3, dtype=jnp.float32).reshape(1, 8, 3) + calls = [] + + class Model: + + def __call__(self, **kwargs): + calls.append(kwargs) + return raw_logits + + batch = diffusion_sft.create_batch_adapter(_config())(_raw_batch()) + actual = diffusion_sft.create_target_aligned_logits_fn(_config())( + Model(), batch.model_inputs + ) + expected = scoring.align_logits_to_targets( + raw_logits, + "shifted", + batch.model_inputs["target_positions"], + batch.model_inputs["target_segmentation"] != 0, + ) + + np.testing.assert_array_equal(actual, expected) + assert calls[0]["enable_dropout"] is False + np.testing.assert_array_equal( + calls[0]["decoder_input_tokens"], _raw_batch()["inputs"] + ) diff --git a/tests/post_training/unit/sft_hooks_test.py b/tests/post_training/unit/sft_hooks_test.py index 9fde5add7e..8b97aceaa5 100644 --- a/tests/post_training/unit/sft_hooks_test.py +++ b/tests/post_training/unit/sft_hooks_test.py @@ -74,6 +74,18 @@ def test_sft_training_hooks_get_total_weights(self): total_weights = training_hooks.get_total_weights(batch) self.assertEqual(total_weights, 3) + def test_sft_training_hooks_prefers_explicit_loss_mask(self): + learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(self.config) + training_hooks = sft_hooks.SFTTrainingHooks(self.config, self.mesh, learning_rate_schedule, goodput_recorder=None) + batch = { + "targets_segmentation": np.ones((2, 4), dtype=np.int32), + "targets_loss_mask": np.array([[1, 0, 0, 0], [0, 1, 1, 0]], dtype=np.int32), + } + + total_weights = training_hooks.get_total_weights(batch) + + self.assertEqual(total_weights, 3) + if __name__ == "__main__": unittest.main() diff --git a/tests/post_training/unit/train_sft_test.py b/tests/post_training/unit/train_sft_test.py index f0df897b3e..f579f8159d 100644 --- a/tests/post_training/unit/train_sft_test.py +++ b/tests/post_training/unit/train_sft_test.py @@ -31,6 +31,7 @@ class TrainSFTTest(unittest.TestCase): def test_validate_config_valid(self): config = SimpleNamespace( optimizer_memory_host_offload=False, + eval_interval=-1, ) # Should not raise any exception train_sft.validate_config(config) @@ -43,6 +44,119 @@ def test_validate_config_invalid_offload(self): with self.assertRaisesRegex(ValueError, "optimizer_memory_host_offload=True is not supported"): train_sft.validate_config(config) + @pytest.mark.cpu_only + def test_validate_config_accepts_weighted_diffusion_accumulation(self): + config = SimpleNamespace( + optimizer_memory_host_offload=False, + training_objective="block_diffusion", + gradient_accumulation_steps=2, + eval_interval=-1, + ) + + train_sft.validate_config(config) + + @pytest.mark.cpu_only + def test_validate_config_accepts_weighted_diffusion_evaluation(self): + config = SimpleNamespace( + optimizer_memory_host_offload=False, + training_objective="block_diffusion", + gradient_accumulation_steps=1, + eval_interval=10, + ) + + train_sft.validate_config(config) + + @pytest.mark.cpu_only + def test_setup_validates_before_initializing_model(self): + config = SimpleNamespace( + optimizer_memory_host_offload=True, + ) + + with ( + mock.patch.object(train_sft, "get_tunix_config") as get_tunix_config, + mock.patch.object(train_sft.model_creation_utils, "from_pretrained") as from_pretrained, + self.assertRaisesRegex(ValueError, "optimizer_memory_host_offload=True"), + ): + train_sft.setup_trainer_state(config) + + get_tunix_config.assert_not_called() + from_pretrained.assert_not_called() + + @pytest.mark.cpu_only + def test_maxtext_loss_wrapper_forwards_diffusion_masks(self): + trainer = mock.MagicMock() + trainer.with_loss_fn.return_value = trainer + config = SimpleNamespace() + train_sft.use_maxtext_loss_function(trainer, config) + wrapped_loss = trainer.with_loss_fn.call_args.args[0] + arrays = { + name: mock.sentinel.__getattr__(name) + for name in ( + "inputs", + "inputs_position", + "inputs_segmentation", + "targets", + "targets_position", + "targets_segmentation", + "completion_mask", + "corruption_mask", + "targets_loss_mask", + ) + } + + with mock.patch.object(train_sft, "loss_fn", return_value=(1.0, {})) as loss_fn_mock: + wrapped_loss(mock.sentinel.model, **arrays) + + forwarded = loss_fn_mock.call_args.args[2] + self.assertEqual(forwarded, arrays) + trainer.with_loss_fn.assert_called_once_with(wrapped_loss, has_aux=True) + + @pytest.mark.cpu_only + def test_diffusion_objective_uses_tunix_adapter(self): + trainer = mock.sentinel.trainer + config = SimpleNamespace(training_objective="block_diffusion") + with ( + mock.patch.object(train_sft.maxtext_diffusion_sft, "create_batch_adapter", return_value="adapter"), + mock.patch.object( + train_sft.maxtext_diffusion_sft, + "create_target_aligned_logits_fn", + return_value="logits_fn", + ), + mock.patch.object( + train_sft.tunix_diffusion_sft, + "configure_diffusion_sft", + return_value=trainer, + ) as configure, + ): + result = train_sft.configure_training_objective(trainer, config) + + self.assertIs(result, trainer) + configure.assert_called_once_with(trainer, "adapter", "logits_fn") + + @pytest.mark.cpu_only + def test_ar_objective_retains_maxtext_trainer_and_loss(self): + config = SimpleNamespace(training_objective="causal_lm") + with ( + mock.patch.object(train_sft, "MaxTextPeftTrainer", return_value=mock.sentinel.trainer) as trainer_type, + mock.patch.object(train_sft, "use_maxtext_loss_function", return_value=mock.sentinel.configured) as configure, + ): + trainer = train_sft._create_trainer("model", "optimizer", "config", config) + result = train_sft.configure_training_objective(trainer, config) + + self.assertIs(trainer, mock.sentinel.trainer) + self.assertIs(result, mock.sentinel.configured) + trainer_type.assert_called_once_with("model", "optimizer", "config") + configure.assert_called_once_with(mock.sentinel.trainer, config) + + @pytest.mark.cpu_only + def test_diffusion_objective_uses_current_tunix_trainer(self): + config = SimpleNamespace(training_objective="block_diffusion") + with mock.patch.object(train_sft.peft_trainer, "PeftTrainer", return_value=mock.sentinel.trainer) as trainer_type: + trainer = train_sft._create_trainer("model", "optimizer", "config", config) + + self.assertIs(trainer, mock.sentinel.trainer) + trainer_type.assert_called_once_with("model", "optimizer", "config") + @pytest.mark.cpu_only def test_train_model_caching_moe(self): """Test that NNX graph caching is disabled for MoE models (num_experts > 1).""" diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index f333748192..32a4bbba0e 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -239,6 +239,107 @@ def test_default_attention_remains_global(self): config = pyconfig.initialize(["", _BASE_CONFIG_PATH, "run_name=test", "steps=1"]) self.assertEqual(config.attention_type, "global") + self.assertEqual(config.training_objective, "causal_lm") + self.assertEqual(config.block_diffusion_mask_id, -1) + + def test_block_diffusion_training_objective_config(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "training_objective=block_diffusion", + "attention=dot_product", + "attention_type=block_diffusion", + "block_diffusion_block_size=7", + "block_diffusion_mask_id=100", + "block_diffusion_min_noise=0.05", + "vocab_size=256", + "max_target_length=2048", + "packing=False", + "dataset_type=hf", + "hf_path=parquet", + "hardware=cpu", + ] + + config = pyconfig.initialize(argv) + + self.assertEqual(config.training_objective, "block_diffusion") + self.assertEqual(config.block_diffusion_mask_id, 100) + self.assertEqual(config.block_diffusion_min_noise, 0.05) + self.assertEqual(config.block_diffusion_logit_alignment, "same_position") + self.assertEqual(config.block_diffusion_canvas_policy, "all_masked") + self.assertNotEqual(config.max_target_length % config.block_diffusion_block_size, 0) + + def test_block_diffusion_training_objective_rejects_incompatible_config(self): + base_overrides = { + "run_name": "test", + "steps": 1, + "training_objective": "block_diffusion", + "attention": "dot_product", + "attention_type": "block_diffusion", + "block_diffusion_block_size": 32, + "block_diffusion_mask_id": 100, + "block_diffusion_min_noise": 0.05, + "vocab_size": 256, + "max_target_length": 2048, + "packing": False, + "dataset_type": "hf", + "hf_path": "parquet", + "hardware": "cpu", + } + cases = [ + ({"attention_type": "global"}, "attention_type='block_diffusion'"), + ({"block_diffusion_mask_id": -1}, "block_diffusion_mask_id"), + ({"block_diffusion_mask_id": 256}, "block_diffusion_mask_id"), + ({"block_diffusion_min_noise": 0.0}, "block_diffusion_min_noise"), + ({"packing": True}, "packing=False"), + ({"mtp_num_layers": 1}, "MTP"), + ({"num_vocab_tiling": 2}, "vocabulary tiling"), + ({"dataset_type": "grain"}, "dataset_type='hf'"), + ({"use_dpo": True}, "DPO"), + ({"use_multimodal": True}, "text-only"), + ({"block_diffusion_logit_alignment": "shifted"}, "seed_and_mask"), + ( + {"block_diffusion_canvas_policy": "seed_and_mask"}, + "same_position/all_masked", + ), + ] + for overrides, expected_regex in cases: + with self.subTest(overrides=overrides): + config = base_overrides | overrides + argv = [ + "", + _BASE_CONFIG_PATH, + *(f"{key}={value}" for key, value in config.items()), + ] + with self.assertRaisesRegex((ValueError, pydantic.ValidationError), expected_regex): + pyconfig.initialize(argv) + + def test_shifted_block_diffusion_requires_seeded_canvas(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "training_objective=block_diffusion", + "attention=dot_product", + "attention_type=block_diffusion", + "block_diffusion_block_size=8", + "block_diffusion_mask_id=100", + "block_diffusion_logit_alignment=shifted", + "block_diffusion_canvas_policy=seed_and_mask", + "vocab_size=256", + "packing=False", + "dataset_type=hf", + "hf_path=parquet", + "hardware=cpu", + ] + + config = pyconfig.initialize(argv) + + self.assertEqual(config.block_diffusion_logit_alignment, "shifted") + self.assertEqual(config.block_diffusion_canvas_policy, "seed_and_mask") @unittest.mock.patch.dict(os.environ, {pyconfig.yaml_key_to_env_key("steps"): "123"}) def test_env_override(self): diff --git a/tests/unit/gradient_accumulation_nnx_test.py b/tests/unit/gradient_accumulation_nnx_test.py index 478fe9a746..5ca56f484d 100644 --- a/tests/unit/gradient_accumulation_nnx_test.py +++ b/tests/unit/gradient_accumulation_nnx_test.py @@ -14,6 +14,8 @@ """Unit tests for the NNX branch of gradient_accumulation_loss_and_grad.""" +# pylint: disable=too-many-positional-arguments + import unittest from dataclasses import dataclass @@ -67,6 +69,20 @@ def _fake_loss_fn(model, config, data, dropout_rng, params, is_train=True): return xent_sum / total_weights, aux +def _zero_weight_loss_fn(model, config, data, dropout_rng, params, is_train=True): + """Produces model-dependent sums with a zero token denominator.""" + del config, dropout_rng, params, is_train + xent_sum = jnp.sum(model(data["inputs"]) ** 2) + 1.0 + aux = { + "xent_sum": xent_sum, + "total_weights": jnp.array(0.0), + "moe_lb_loss": jnp.array(0.0), + "indexer_loss": jnp.array(0.0), + "mtp_loss": jnp.array(0.0), + } + return xent_sum, aux + + class TestGradientAccumulationNNX(unittest.TestCase): """Cover the NNX path of gradient_accumulation_loss_and_grad.""" @@ -151,6 +167,22 @@ def test_nnx_with_shard_optimizer_over_data_casts_to_bf16(self): for g in jax.tree.leaves(raw_grads): self.assertTrue(jnp.all(jnp.isfinite(g))) + def test_zero_total_weights_returns_zero_loss_and_gradients(self): + loss, _, raw_grads = gradient_accumulation.gradient_accumulation_loss_and_grad( + _zero_weight_loss_fn, + self.cfg, + self.model, + params=None, + params_shardings=self._params_shardings(), + data=self.data, + dropout_rng=None, + ) + + self.assertEqual(float(loss), 0.0) + for gradient in jax.tree.leaves(raw_grads): + self.assertTrue(jnp.all(jnp.isfinite(gradient))) + self.assertTrue(jnp.all(gradient == 0)) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/hf_data_processing_test.py b/tests/unit/hf_data_processing_test.py index 93564514a3..a23525af02 100644 --- a/tests/unit/hf_data_processing_test.py +++ b/tests/unit/hf_data_processing_test.py @@ -15,6 +15,7 @@ """Tests for Hugging Face data processing.""" import sys +from types import SimpleNamespace import unittest import os.path @@ -26,6 +27,7 @@ from maxtext.configs import pyconfig from maxtext.input_pipeline import hf_data_processing from maxtext.input_pipeline import input_pipeline_interface +from maxtext.input_pipeline import input_pipeline_utils from maxtext.common.gcloud_stub import is_decoupled from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory @@ -126,5 +128,75 @@ def get_first_batch(iterator): self.assertTrue((train_batch1["targets"] == train_batch2["targets"]).all()) # pytype: disable=unsupported-operands +@pytest.mark.cpu_only +class TrainingObjectiveTransformTest(unittest.TestCase): + """Tests objective selection at the HF pipeline boundary.""" + + def test_default_objective_keeps_shift_data(self): + transform = hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + SimpleNamespace(), + shift=True, + use_dpo=False, + use_sft=False, + completion_only=False, + packing=False, + pad_id=0, + bos_token_id=1, + ) + + self.assertIsInstance(transform, input_pipeline_utils.ShiftData) + self.assertEqual(transform.ignored_ids, [0, 1]) + + def test_explicit_block_diffusion_objective_replaces_shift(self): + config = SimpleNamespace( + training_objective="block_diffusion", + block_diffusion_block_size=4, + block_diffusion_mask_id=99, + block_diffusion_min_noise=0.05, + block_diffusion_canvas_policy="seed_and_mask", + block_diffusion_logit_alignment="shifted", + ) + + transform = hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + config, + shift=True, + use_dpo=False, + use_sft=True, + completion_only=True, + packing=False, + pad_id=0, + bos_token_id=1, + ) + + self.assertIsInstance(transform, input_pipeline_utils.BlockDiffusionCorruption) + self.assertEqual(transform.block_size, 4) + self.assertEqual(transform.mask_id, 99) + self.assertEqual(transform.min_noise, 0.05) + self.assertTrue(transform.completion_only) + self.assertTrue(transform.seed_first_token) + self.assertTrue(transform.include_seed_in_loss) + + def test_block_diffusion_rejects_packing(self): + config = SimpleNamespace( + training_objective="block_diffusion", + block_diffusion_block_size=4, + block_diffusion_mask_id=99, + block_diffusion_min_noise=0.05, + block_diffusion_canvas_policy="all_masked", + ) + + with self.assertRaisesRegex(ValueError, "packing=False"): + hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + config, + shift=True, + use_dpo=False, + use_sft=False, + completion_only=False, + packing=True, + pad_id=0, + bos_token_id=1, + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/input_pipeline_utils_test.py b/tests/unit/input_pipeline_utils_test.py index 0d5adae4a9..fe5ce30cda 100644 --- a/tests/unit/input_pipeline_utils_test.py +++ b/tests/unit/input_pipeline_utils_test.py @@ -17,7 +17,14 @@ import pytest import unittest -from maxtext.input_pipeline.input_pipeline_utils import compute_file_sharding +import numpy as np + +from maxtext.input_pipeline.input_pipeline_utils import ( + BlockDiffusionCorruption, + PadOrTrimToMaxLength, + SFTPromptMasking, + compute_file_sharding, +) @pytest.mark.cpu_only @@ -109,5 +116,210 @@ def test_no_row_shard_when_only_one_reader(self): self.assertIsNone(row_shard) +def _make_block_diffusion_batch(): + """Builds a padded, completion-only SFT batch with a partial final block.""" + inputs = np.asarray( + [ + [11, 12, 21, 22, 23, 24, 25, 26, 27, 0], + [31, 41, 42, 43, 44, 45, 0, 0, 0, 0], + ], + dtype=np.int32, + ) + targets = np.asarray( + [ + [11, 12, 21, 22, 23, 24, 25, 26, 27, 0], + [31, 41, 42, 43, 44, 45, 0, 0, 0, 0], + ], + dtype=np.int32, + ) + completion_mask = np.asarray( + [ + [0, 0, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 1, 1, 1, 1, 1, 0, 0, 0, 0], + ], + dtype=np.int32, + ) + positions = np.broadcast_to(np.arange(inputs.shape[1], dtype=np.int32), inputs.shape).copy() + return { + "inputs": inputs, + "targets": targets, + "inputs_segmentation": (inputs != 0).astype(np.int32), + "targets_segmentation": (targets != 0).astype(np.int32), + "inputs_position": positions.copy(), + "targets_position": positions.copy(), + "completion_mask": completion_mask, + } + + +class _ControlledRng: + """Supplies distinct noise rates while making token draws deterministic.""" + + def __init__(self): + self.uniform_calls = 0 + self.random_calls = 0 + + def uniform(self, low, high, size): + del low, high + self.uniform_calls += 1 + return np.asarray([[[0.1], [0.9], [0.1]], [[0.1], [0.9], [0.1]]]).reshape(size) + + def random(self, size): + self.random_calls += 1 + return np.full(size, 0.5) + + +@pytest.mark.cpu_only +class BlockDiffusionCorruptionTest(unittest.TestCase): + """Tests aligned block-diffusion corruption metadata.""" + + def _apply(self, seed=0): + transform = BlockDiffusionCorruption(block_size=4, mask_id=99, min_noise=1.0e-3, completion_only=True) + return transform.random_map(_make_block_diffusion_batch(), np.random.default_rng(seed)) + + def test_preserves_targets_and_segmentation(self): + clean = _make_block_diffusion_batch() + output = self._apply() + + for key in ( + "targets", + "inputs_segmentation", + "targets_segmentation", + "inputs_position", + "targets_position", + ): + np.testing.assert_array_equal(output[key], clean[key]) + np.testing.assert_array_equal(output["completion_mask"], clean["completion_mask"]) + self.assertTrue(np.all(output["targets_loss_mask"] <= output["completion_mask"])) + self.assertFalse(np.shares_memory(output["completion_mask"], output["targets_loss_mask"])) + np.testing.assert_array_equal(output["corruption_mask"], output["targets_loss_mask"]) + + loss_positions = output["targets_loss_mask"] != 0 + self.assertTrue(np.all(output["inputs"][loss_positions] == 99)) + np.testing.assert_array_equal(output["inputs"][~loss_positions], clean["inputs"][~loss_positions]) + + def test_every_eligible_full_or_partial_block_has_loss(self): + output = self._apply(seed=17) + completion_mask = output["completion_mask"] != 0 + loss_mask = output["targets_loss_mask"] != 0 + + for row in range(completion_mask.shape[0]): + for start in range(0, completion_mask.shape[1], 4): + stop = min(start + 4, completion_mask.shape[1]) + if np.any(completion_mask[row, start:stop]): + self.assertTrue(np.any(loss_mask[row, start:stop]), (row, start, stop)) + self.assertTrue(loss_mask[0, 8]) + + def test_corruption_rate_is_sampled_per_block(self): + rng = _ControlledRng() + transform = BlockDiffusionCorruption(block_size=4, mask_id=99, completion_only=True) + + output = transform.random_map(_make_block_diffusion_batch(), rng) + + self.assertEqual(rng.uniform_calls, 1) + self.assertEqual(rng.random_calls, 2) + self.assertEqual(int(output["targets_loss_mask"][0, :4].sum()), 1) + self.assertEqual(int(output["targets_loss_mask"][0, 4:8].sum()), 4) + + def test_deterministic_for_fixed_seed(self): + output_a = self._apply(seed=123) + output_b = self._apply(seed=123) + for key, value in output_a.items(): + np.testing.assert_array_equal(value, output_b[key]) + + def test_invalid_parameters_raise(self): + with self.assertRaises(ValueError): + BlockDiffusionCorruption(block_size=0, mask_id=99) + with self.assertRaises(ValueError): + BlockDiffusionCorruption(block_size=4, mask_id=99, min_noise=0.0) + with self.assertRaisesRegex(ValueError, "requires seed_first_token"): + BlockDiffusionCorruption(block_size=4, mask_id=99, include_seed_in_loss=True) + + def test_completion_only_requires_explicit_role_mask(self): + clean = _make_block_diffusion_batch() + del clean["completion_mask"] + + with self.assertRaisesRegex(ValueError, "explicit completion_mask"): + self._apply_transform(clean) + + def test_completion_mask_does_not_replace_full_validity(self): + clean = _make_block_diffusion_batch() + transform = BlockDiffusionCorruption(block_size=4, mask_id=99, min_noise=1.0, completion_only=False) + + output = transform.random_map(clean, np.random.default_rng(0)) + + np.testing.assert_array_equal(output["targets_segmentation"], clean["targets_segmentation"]) + np.testing.assert_array_equal(output["completion_mask"], clean["completion_mask"]) + self.assertTrue(np.all(output["targets_loss_mask"][clean["inputs_segmentation"] != 0])) + self.assertTrue(np.any(output["targets_loss_mask"] > output["completion_mask"])) + + def test_all_padding_has_finite_zero_masks(self): + clean = _make_block_diffusion_batch() + for key in ( + "inputs", + "targets", + "inputs_segmentation", + "targets_segmentation", + "completion_mask", + ): + clean[key] = np.zeros_like(clean[key]) + + output = self._apply_transform(clean) + + self.assertFalse(output["corruption_mask"].any()) + self.assertFalse(output["targets_loss_mask"].any()) + + def test_shifted_seed_canvas_keeps_anchors_clean_but_supervised(self): + clean = _make_block_diffusion_batch() + transform = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + min_noise=1.0, + completion_only=False, + seed_first_token=True, + include_seed_in_loss=True, + ) + + output = transform.random_map(clean, np.random.default_rng(0)) + + self.assertFalse(output["corruption_mask"][:, ::4].any()) + expected_loss = clean["targets_segmentation"].copy() + expected_loss[:, 0] = 0 + np.testing.assert_array_equal(output["targets_loss_mask"], expected_loss) + self.assertTrue(output["targets_loss_mask"][0, 4]) + self.assertTrue(output["targets_loss_mask"][0, 8]) + + def _apply_transform(self, batch): + transform = BlockDiffusionCorruption(block_size=4, mask_id=99, completion_only=True) + return transform.random_map(batch, np.random.default_rng(0)) + + +@pytest.mark.cpu_only +class BlockDiffusionSftMaskTest(unittest.TestCase): + """Tests role metadata before randomized corruption.""" + + def test_target_aligned_sft_preserves_clean_tokens_and_completion_roles(self): + transform = SFTPromptMasking( + text_column_name="messages", + completion_only=True, + max_target_length=8, + unk_id=0, + target_aligned=True, + emit_completion_mask=True, + ) + element = { + "messages": [[11, 12], [21, 22], [31], [41, 42]], + "is_prompt": [True, False, True, False], + } + + output = transform.map(element) + padded = PadOrTrimToMaxLength(max_length=8, pad_id=0).map(output) + + np.testing.assert_array_equal(padded["targets"], padded["inputs"]) + np.testing.assert_array_equal(padded["completion_mask"], [0, 0, 1, 1, 0, 1, 1, 0]) + np.testing.assert_array_equal(padded["targets_segmentation"], [1, 1, 1, 1, 1, 1, 1, 0]) + self.assertNotIn("completion_mask_segmentation", padded) + self.assertNotIn("completion_mask_position", padded) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/max_utils_test.py b/tests/unit/max_utils_test.py index 097d85d491..7dcc6cea83 100644 --- a/tests/unit/max_utils_test.py +++ b/tests/unit/max_utils_test.py @@ -23,10 +23,12 @@ import jax from jax import numpy as jnp from jax import random +from maxtext.common.common_types import ReorderStrategy from maxtext.configs import pyconfig from maxtext.utils import max_utils from maxtext.utils.train_utils import setup_train_loop from tests.utils.test_helpers import get_test_config_path +import numpy as np import optax import pytest @@ -655,6 +657,31 @@ def test_reorder_roundtrip(self): # 3. Assert roundtrip is lossless self.assertTrue(jnp.allclose(x, restored, rtol=1e-5, atol=1e-6)) + def test_block_diffusion_masks_are_reordered_with_tokens(self): + completion_mask = jnp.arange(16, dtype=jnp.int32).reshape(1, 16) + targets_loss_mask = completion_mask + 100 + corruption_mask = completion_mask + 50 + batch = { + "inputs": completion_mask + 200, + "completion_mask": completion_mask, + "corruption_mask": corruption_mask, + "targets_loss_mask": targets_loss_mask, + } + + reordered = max_utils.reorder_causal_load_balanced( + batch, + cp_size=2, + reorder_strategy=ReorderStrategy.DUAL_CHUNK_SWAP, + hardware="cpu", + ) + + np.testing.assert_array_equal(reordered["completion_mask"], max_utils.reorder_sequence(completion_mask, 2)) + np.testing.assert_array_equal(reordered["corruption_mask"], max_utils.reorder_sequence(corruption_mask, 2)) + np.testing.assert_array_equal( + reordered["targets_loss_mask"], + max_utils.reorder_sequence(targets_loss_mask, 2), + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index 28c9f6572a..7f33735c1e 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/tests/unit/maxtext_utils_test.py @@ -1140,13 +1140,22 @@ def test_linen_in_shardings_includes_rng(self): class TestGetShapedBatch(unittest.TestCase): """Tests for get_shaped_batch.""" - def _make_cfg(self, *, enable_diloco=False, use_multimodal=False, use_audio=False): + def _make_cfg( + self, + *, + enable_diloco=False, + use_multimodal=False, + use_audio=False, + training_objective="causal_lm", + ): + """Builds the config subset consumed by get_shaped_batch.""" cfg = MagicMock() cfg.enable_diloco = enable_diloco cfg.global_batch_size_to_load = 4 cfg.max_target_length = 16 cfg.use_multimodal = use_multimodal cfg.use_audio = use_audio + cfg.training_objective = training_objective if enable_diloco: cfg.num_diloco_replicas = 2 return cfg @@ -1169,6 +1178,23 @@ def test_standard_shape(self): expected_shape = (cfg.global_batch_size_to_load, cfg.max_target_length) self.assertEqual(batch["inputs"].shape, expected_shape) + def test_block_diffusion_masks_are_in_shaped_batch(self): + cfg = self._make_cfg(training_objective="block_diffusion") + + batch = maxtext_utils.get_shaped_batch(cfg) + + self.assertEqual(batch["completion_mask"].shape, batch["inputs"].shape) + self.assertEqual(batch["corruption_mask"].shape, batch["inputs"].shape) + self.assertEqual(batch["targets_loss_mask"].shape, batch["inputs"].shape) + self.assertEqual(batch["targets_loss_mask"].dtype, jnp.int32) + + def test_causal_shaped_batch_has_no_diffusion_masks(self): + batch = maxtext_utils.get_shaped_batch(self._make_cfg()) + + self.assertNotIn("completion_mask", batch) + self.assertNotIn("corruption_mask", batch) + self.assertNotIn("targets_loss_mask", batch) + def test_diloco_shape(self): cfg = self._make_cfg(enable_diloco=True) batch = maxtext_utils.get_shaped_batch(cfg) diff --git a/tests/unit/pre_train_loss_mask_test.py b/tests/unit/pre_train_loss_mask_test.py new file mode 100644 index 0000000000..07e4f6200b --- /dev/null +++ b/tests/unit/pre_train_loss_mask_test.py @@ -0,0 +1,264 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests explicit target-loss masking in the pre-training loss.""" + +# pylint: disable=too-many-positional-arguments + +from dataclasses import dataclass +import unittest +from unittest import mock + +from flax import linen as nn +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.trainers.pre_train import train as pre_train +from maxtext.diffusion import scoring as diffusion_scoring + + +@dataclass +class _Config: + """Configuration subset consumed by pre-training loss_fn.""" + + micro_batch_size_to_train_on: int = 2 + micro_batch_size_to_eval_on: int = 2 + vocab_size: int = 8 + z_loss_multiplier: float = 0.0 + enable_dropout: bool = False + use_multimodal: bool = False + use_indexer: bool = False + indexer_sparse_training: bool = False + indexer_loss_scaling_factor: float = 0.0 + num_vocab_tiling: int = 1 + num_experts: int = 1 + routed_bias: bool = False + routed_bias_update_rate: float = 0.0 + mtp_num_layers: int = 0 + mtp_eval_target_module: int = 0 + use_qk_clip: bool = False + use_tunix_gradient_accumulation: bool = False + gradient_accumulation_steps: int = 1 + shard_mode: int = 0 + debug_sharding: bool = False + weight_sparsity_n: int = 0 + weight_sparsity_m: int = 0 + attention_type: str = "global" + training_objective: str = "causal_lm" + block_diffusion_logit_alignment: str = "same_position" + block_diffusion_canvas_policy: str = "all_masked" + block_diffusion_block_size: int = 4 + + +class _UniformNnxDecoder(nnx.Module): + """Returns uniform logits through the NNX call contract.""" + + def __init__(self, vocab_size): + self.vocab_size = vocab_size + self.mesh = jax.make_mesh((1, 1, 1, 1), ("data", "fsdp", "expert", "context")) + + def __call__( + self, + decoder_input_tokens, + decoder_positions, + decoder_segment_ids=None, + encoder_images=None, + encoder_image_masks=None, + enable_dropout=False, + decoder_target_tokens=None, + decoder_target_mask=None, + ): + del decoder_positions, decoder_segment_ids, encoder_images, encoder_image_masks + del enable_dropout, decoder_target_tokens, decoder_target_mask + return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32) + + +class _UniformLinenDecoder(nn.Module): + """Returns uniform logits through the Linen call contract.""" + + vocab_size: int + mesh: object + + @nn.compact + def __call__( + self, + decoder_input_tokens, + decoder_positions, + decoder_segment_ids=None, + encoder_images=None, + encoder_image_masks=None, + enable_dropout=False, + decoder_target_tokens=None, + decoder_target_mask=None, + ): + del decoder_positions, decoder_segment_ids, encoder_images, encoder_image_masks + del enable_dropout, decoder_target_tokens, decoder_target_mask + return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32) + + +def _make_data(include_loss_mask=True): + """Builds a batch whose explicit loss mask differs from segmentation.""" + data = { + "inputs": jnp.zeros((2, 4), dtype=jnp.int32), + "inputs_position": jnp.broadcast_to(jnp.arange(4), (2, 4)), + "inputs_segmentation": jnp.ones((2, 4), dtype=jnp.int32), + "targets": jnp.zeros((2, 4), dtype=jnp.int32), + "targets_segmentation": jnp.asarray([[1, 1, 1, 1], [1, 1, 1, 0]], dtype=jnp.int32), + "completion_mask": jnp.asarray([[1, 1, 1, 1], [1, 1, 1, 0]], dtype=jnp.int32), + "corruption_mask": jnp.asarray([[1, 0, 1, 0], [0, 1, 0, 0]], dtype=jnp.int32), + } + if include_loss_mask: + data["targets_loss_mask"] = jnp.asarray([[1, 0, 1, 0], [0, 1, 0, 0]], dtype=jnp.int32) + return data + + +class PreTrainLossMaskTest(unittest.TestCase): + """Checks explicit masks in both Linen and NNX loss branches.""" + + def setUp(self): + super().setUp() + self.config = _Config() + self.per_token_xent = jnp.arange(1, 9, dtype=jnp.float32).reshape(2, 4) + self.per_token_z_loss = self.per_token_xent / 10.0 + + def _cross_entropy_patch(self): + return mock.patch.object( + pre_train.max_utils, + "cross_entropy_with_logits", + return_value=(self.per_token_xent, self.per_token_z_loss), + ) + + def _use_block_diffusion(self): + self.config.attention_type = "block_diffusion" + self.config.training_objective = "block_diffusion" + + def _assert_explicit_mask_result(self, loss, aux): + expected_mask = _make_data()["targets_loss_mask"] != 0 + expected_xent = jnp.sum(self.per_token_xent * expected_mask) + expected_z_loss = jnp.sum(self.per_token_z_loss * expected_mask) / 3.0 + self.assertEqual(int(aux["total_weights"]), 3) + self.assertAlmostEqual(float(aux["xent_sum"]), float(expected_xent)) + self.assertAlmostEqual(float(loss), float(expected_xent / 3.0)) + self.assertAlmostEqual(float(aux["z_loss"]), float(expected_z_loss)) + + def test_nnx_loss_uses_targets_loss_mask(self): + self._use_block_diffusion() + model = _UniformNnxDecoder(self.config.vocab_size) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, _make_data(), None, None, is_train=True) + + self._assert_explicit_mask_result(loss, aux) + + def test_linen_loss_uses_targets_loss_mask(self): + self._use_block_diffusion() + mesh = jax.make_mesh((1, 1, 1, 1), ("data", "fsdp", "expert", "context")) + model = _UniformLinenDecoder(vocab_size=self.config.vocab_size, mesh=mesh) + data = _make_data() + variables = model.init( + jax.random.key(0), + data["inputs"], + data["inputs_position"], + decoder_segment_ids=data["inputs_segmentation"], + decoder_target_tokens=data["targets"], + decoder_target_mask=data["targets_segmentation"], + ) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, data, jax.random.key(1), variables, is_train=True) + + self._assert_explicit_mask_result(loss, aux) + + def test_causal_fallback_uses_targets_segmentation(self): + data = _make_data(include_loss_mask=False) + model = _UniformNnxDecoder(self.config.vocab_size) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + expected_mask = data["targets_segmentation"] != 0 + expected_xent = jnp.sum(self.per_token_xent * expected_mask) + self.assertEqual(int(aux["total_weights"]), 7) + self.assertAlmostEqual(float(aux["xent_sum"]), float(expected_xent)) + self.assertAlmostEqual(float(loss), float(expected_xent / 7.0)) + + def test_zero_explicit_mask_has_finite_zero_loss(self): + self._use_block_diffusion() + data = _make_data() + data["targets_loss_mask"] = jnp.zeros_like(data["targets_loss_mask"]) + model = _UniformNnxDecoder(self.config.vocab_size) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + self.assertEqual(int(aux["total_weights"]), 0) + self.assertEqual(float(aux["xent_sum"]), 0.0) + self.assertTrue(bool(jnp.isfinite(loss))) + self.assertEqual(float(loss), 0.0) + + def test_causal_loss_rejects_block_bidirectional_attention(self): + self.config.attention_type = "block_diffusion" + model = _UniformNnxDecoder(self.config.vocab_size) + + with self.assertRaisesRegex(ValueError, "would leak"): + pre_train.loss_fn(model, self.config, _make_data(), None, None, is_train=True) + + def test_block_diffusion_requires_all_explicit_masks(self): + self._use_block_diffusion() + model = _UniformNnxDecoder(self.config.vocab_size) + + for mask_name in ("completion_mask", "corruption_mask", "targets_loss_mask"): + with self.subTest(mask_name=mask_name): + data = _make_data() + del data[mask_name] + with self.assertRaisesRegex(ValueError, mask_name): + pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + def test_logit_alignment_matches_runtime_contract(self): + logits = jnp.arange(8, dtype=jnp.float32).reshape(1, 4, 2) + + same_position = diffusion_scoring.align_logits_to_targets(logits, "same_position") + shifted = diffusion_scoring.align_logits_to_targets(logits, "shifted") + + np.testing.assert_array_equal(same_position, logits) + np.testing.assert_array_equal(shifted, logits[:, [0, 0, 1, 2], :]) + + def test_shifted_alignment_uses_logical_positions_after_reordering(self): + positions = jnp.asarray([[0, 3, 1, 2]], dtype=jnp.int32) + logits = positions[..., None].astype(jnp.float32) + + shifted = diffusion_scoring.align_logits_to_targets( + logits, + "shifted", + positions=positions, + validity_mask=jnp.ones_like(positions, dtype=jnp.bool_), + ) + + np.testing.assert_array_equal(shifted[..., 0], [[0, 2, 0, 1]]) + + def test_seeded_shifted_canvas_excludes_only_sequence_origin(self): + self.config.attention_type = "block_diffusion" + self.config.training_objective = "block_diffusion" + self.config.block_diffusion_logit_alignment = "shifted" + self.config.block_diffusion_canvas_policy = "seed_and_mask" + data = _make_data() + data["targets_loss_mask"] = jnp.ones_like(data["targets_segmentation"]) + model = _UniformNnxDecoder(self.config.vocab_size) + + with self._cross_entropy_patch(): + _, aux = pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + self.assertEqual(int(aux["total_weights"]), 5) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/synthetic_data_test.py b/tests/unit/synthetic_data_test.py index 9e29c04a15..f713037a13 100644 --- a/tests/unit/synthetic_data_test.py +++ b/tests/unit/synthetic_data_test.py @@ -15,6 +15,7 @@ """Tests for synthetic data sharding.""" import sys +from types import SimpleNamespace from absl.testing import absltest from absl.testing import parameterized import jax @@ -22,7 +23,7 @@ from jax.experimental import mesh_utils from maxtext.configs import pyconfig -from maxtext.input_pipeline.synthetic_data_processing import SyntheticDataIterator +from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator, SyntheticDataIterator from maxtext.utils import sharding as maxtext_sharding from tests.utils.test_helpers import get_test_config_path @@ -75,5 +76,25 @@ def test_synthetic_data_sharding(self, mesh_shape, shard_mode): self.assertEqual(inputs.sharding, expected_sharding) +class PlaceHolderDataTest(absltest.TestCase): + """Tests placeholder batch metadata for non-loading hosts.""" + + def test_block_diffusion_placeholder_has_zero_loss_masks(self): + config = SimpleNamespace( + global_batch_size_to_load=jax.process_count(), + max_target_length=8, + training_objective="block_diffusion", + ) + + batch = next(PlaceHolderDataIterator.get_place_holder_synthetic_data(config)) + + self.assertIn("completion_mask", batch) + self.assertIn("corruption_mask", batch) + self.assertIn("targets_loss_mask", batch) + self.assertFalse(batch["completion_mask"].any()) + self.assertFalse(batch["corruption_mask"].any()) + self.assertFalse(batch["targets_loss_mask"].any()) + + if __name__ == "__main__": absltest.main() From 7f18a02f805f4e8a82376e50fdbde8247c5740df Mon Sep 17 00:00:00 2001 From: ethannnnnn Date: Thu, 23 Jul 2026 00:01:42 -0700 Subject: [PATCH 3/6] [maxtext] Add reusable block-diffusion rollout transitions Add a model-independent low-confidence block rollout that consumes target- aligned logits. It supports the two public model contracts, logical-position block boundaries after context-parallel reordering, heterogeneous batches, partial final blocks, confidence-threshold commits, and forced-argmax progress. The initial OPD scope validates a single contiguous completion suffix so clean future turns cannot leak through bidirectional intra-block attention. Shifted rollouts also require logical position zero to remain prompt context. Test Plan: - 7 passed in tests/unit/diffusion_denoise_test.py, including jax.jit execution - Pylint 10.00/10 - Pyink, pycompile, and git diff checks pass --- src/maxtext/diffusion/denoise.py | 208 +++++++++++++++++++++ tests/unit/diffusion_denoise_test.py | 267 +++++++++++++++++++++++++++ 2 files changed, 475 insertions(+) create mode 100644 src/maxtext/diffusion/denoise.py create mode 100644 tests/unit/diffusion_denoise_test.py diff --git a/src/maxtext/diffusion/denoise.py b/src/maxtext/diffusion/denoise.py new file mode 100644 index 0000000000..2f99b91069 --- /dev/null +++ b/src/maxtext/diffusion/denoise.py @@ -0,0 +1,208 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Model-independent block-diffusion rollout state transitions.""" + +from collections.abc import Callable + +import jax +import jax.numpy as jnp +import numpy as np + + +def _validate_shapes(initial_tokens, positions, validity_mask, completion_mask): + """Checks that all token-level rollout arrays share a batch-major shape.""" + expected_shape = tuple(initial_tokens.shape) + if len(expected_shape) != 2: + raise ValueError(f"initial_tokens must have shape [batch, length]; received {expected_shape}") + for name, value in ( + ("positions", positions), + ("validity_mask", validity_mask), + ("completion_mask", completion_mask), + ): + if tuple(value.shape) != expected_shape: + raise ValueError(f"{name} must match initial_tokens shape; received {tuple(value.shape)} and {expected_shape}") + + +def _concrete_numpy(value): + if isinstance(value, jax.core.Tracer): + return None + if isinstance(value, jax.Array) and not value.is_fully_addressable: + return None + return np.asarray(value) + + +def _validate_logical_positions(positions, validity_mask, completion_mask, *, shifted_seed): + """Checks logical sequence invariants on eager, host-addressable inputs.""" + concrete_positions = _concrete_numpy(positions) + concrete_validity = _concrete_numpy(validity_mask) + concrete_completion = _concrete_numpy(completion_mask) + if concrete_positions is None or concrete_validity is None or concrete_completion is None: + return + concrete_validity = np.asarray(concrete_validity, dtype=bool) + concrete_completion = np.asarray(concrete_completion, dtype=bool) & concrete_validity + sequence_length = concrete_positions.shape[1] + for row in range(concrete_positions.shape[0]): + valid_positions = np.asarray(concrete_positions[row, concrete_validity[row]]) + expected_positions = np.arange(valid_positions.size, dtype=valid_positions.dtype) + if valid_positions.size and not np.array_equal(np.sort(valid_positions), expected_positions): + raise ValueError( + "valid logical positions must be unique and contiguous from zero within the physical sequence length" + ) + if np.any(valid_positions < 0) or np.any(valid_positions >= sequence_length): + raise ValueError("valid logical positions must be nonnegative and smaller than the physical sequence length") + if shifted_seed: + position_zero = concrete_validity[row] & (concrete_positions[row] == 0) + if np.count_nonzero(position_zero) != 1 or np.any(concrete_completion[row] & position_zero): + raise ValueError("shifted seed-and-mask rollout requires exactly one prompt token at logical position zero") + + +def validate_completion_suffix(positions, validity_mask, completion_mask, *, shifted_seed=False): + """Validates the initial OPD rollout scope when eager values are available. + + The first integration intentionally supports one prompt followed by one + completion per sequence. This prevents clean future turns from appearing in + the bidirectional attention block of an earlier generated completion. + """ + _validate_logical_positions(positions, validity_mask, completion_mask, shifted_seed=shifted_seed) + concrete_positions = _concrete_numpy(positions) + concrete_validity = _concrete_numpy(validity_mask) + concrete_completion = _concrete_numpy(completion_mask) + if concrete_positions is None or concrete_validity is None or concrete_completion is None: + return + concrete_validity = np.asarray(concrete_validity, dtype=bool) + concrete_completion = np.asarray(concrete_completion, dtype=bool) & concrete_validity + for row in range(concrete_positions.shape[0]): + valid_indices = np.flatnonzero(concrete_validity[row]) + if valid_indices.size == 0: + continue + ordered_indices = valid_indices[np.argsort(concrete_positions[row, valid_indices])] + ordered_completion = concrete_completion[row, ordered_indices] + completion_indices = np.flatnonzero(ordered_completion) + if completion_indices.size and not np.all(ordered_completion[completion_indices[0] :]): + raise ValueError("diffusion OPD rollout requires completion_mask to be a contiguous suffix") + + +def low_confidence_generate( + logits_fn: Callable[[jax.Array], jax.Array], + initial_tokens: jax.Array, + positions: jax.Array, + validity_mask: jax.Array, + completion_mask: jax.Array, + *, + block_size: int, + mask_id: int, + logit_alignment: str, + canvas_policy: str, + confidence_threshold: float = 0.9, + temperature: float = 1.0, + max_denoise_steps: int | None = None, +) -> jax.Array: + """Generates a completion block by block with confidence-based commits. + + ``logits_fn`` must return target-aligned logits for the current token canvas. + Each denoising step commits every token at or above the confidence threshold; + if a row has no such token, its highest-confidence unresolved token is forced + to commit. This guarantees progress for heterogeneous and partial blocks. + """ + _validate_shapes(initial_tokens, positions, validity_mask, completion_mask) + valid_contracts = {("same_position", "all_masked"), ("shifted", "seed_and_mask")} + if (logit_alignment, canvas_policy) not in valid_contracts: + raise ValueError( + "rollout supports only same_position/all_masked or shifted/seed_and_mask; " + f"received {logit_alignment}/{canvas_policy}" + ) + if block_size <= 0: + raise ValueError(f"block_size must be positive; received {block_size}") + if not 0.0 <= confidence_threshold <= 1.0: + raise ValueError(f"confidence_threshold must be in [0, 1]; received {confidence_threshold}") + if temperature <= 0.0: + raise ValueError(f"temperature must be positive; received {temperature}") + if max_denoise_steps is None: + max_denoise_steps = block_size + if max_denoise_steps < block_size: + raise ValueError( + f"max_denoise_steps must be at least block_size ({block_size}) to guarantee completion; " + f"received {max_denoise_steps}" + ) + + shifted_seed = logit_alignment == "shifted" + validate_completion_suffix(positions, validity_mask, completion_mask, shifted_seed=shifted_seed) + validity_mask = jnp.asarray(validity_mask, dtype=jnp.bool_) + completion_mask = jnp.asarray(completion_mask, dtype=jnp.bool_) & validity_mask + positions = jnp.asarray(positions, dtype=jnp.int32) + canvas = jnp.where(completion_mask, jnp.asarray(mask_id, initial_tokens.dtype), initial_tokens) + block_ids = positions // block_size + num_blocks = (initial_tokens.shape[1] + block_size - 1) // block_size + + def propose(current_canvas): + logits = logits_fn(current_canvas) + expected_prefix = tuple(current_canvas.shape) + if len(logits.shape) != 3 or tuple(logits.shape[:2]) != expected_prefix: + raise ValueError( + "logits_fn must return [batch, length, vocab] target-aligned logits; " + f"received {tuple(logits.shape)} for canvas {expected_prefix}" + ) + vocab_size = logits.shape[-1] + if vocab_size < 2: + raise ValueError("logits_fn must expose at least two vocabulary entries so the mask token can be excluded") + if not 0 <= mask_id < vocab_size: + raise ValueError(f"mask_id must satisfy 0 <= mask_id < vocab_size ({vocab_size}); received {mask_id}") + scaled_logits = jnp.asarray(logits, dtype=jnp.float32) / temperature + scaled_logits = scaled_logits.at[..., mask_id].set(-jnp.inf) + probabilities = jax.nn.softmax(scaled_logits, axis=-1) + return jnp.argmax(scaled_logits, axis=-1).astype(initial_tokens.dtype), jnp.max(probabilities, axis=-1) + + def generate_block(block_id, current_canvas): + in_block = completion_mask & (block_ids == block_id) + + def run_active_block(active_canvas): + anchors = in_block & (positions % block_size == 0) if shifted_seed else jnp.zeros_like(in_block) + + def generate_anchors(anchor_canvas): + anchor_tokens, _ = propose(anchor_canvas) + return jnp.where(anchors, anchor_tokens, anchor_canvas) + + if shifted_seed: + active_canvas = jax.lax.cond(jnp.any(anchors), generate_anchors, lambda value: value, active_canvas) + unresolved = in_block & ~anchors + + def continue_denoising(state): + step, _, remaining = state + return (step < max_denoise_steps) & jnp.any(remaining) + + def denoise_step(state): + step, step_canvas, remaining = state + proposed_tokens, confidence = propose(step_canvas) + commits = remaining & (confidence >= confidence_threshold) + row_needs_fallback = jnp.any(remaining, axis=1) & ~jnp.any(commits, axis=1) + fallback_confidence = jnp.max(jnp.where(remaining, confidence, -jnp.inf), axis=1) + tied_for_fallback = remaining & (confidence == fallback_confidence[:, None]) + fallback_positions = jnp.where(tied_for_fallback, positions, positions.shape[1]) + fallback_indices = jnp.argmin(fallback_positions, axis=1) + fallback = jax.nn.one_hot(fallback_indices, remaining.shape[1], dtype=jnp.bool_) + commits |= fallback & row_needs_fallback[:, None] + step_canvas = jnp.where(commits, proposed_tokens, step_canvas) + return step + 1, step_canvas, remaining & ~commits + + _, active_canvas, _ = jax.lax.while_loop( + continue_denoising, + denoise_step, + (jnp.asarray(0, dtype=jnp.int32), active_canvas, unresolved), + ) + return active_canvas + + return jax.lax.cond(jnp.any(in_block), run_active_block, lambda value: value, current_canvas) + + return jax.lax.fori_loop(0, num_blocks, generate_block, canvas) diff --git a/tests/unit/diffusion_denoise_test.py b/tests/unit/diffusion_denoise_test.py new file mode 100644 index 0000000000..37bf84491d --- /dev/null +++ b/tests/unit/diffusion_denoise_test.py @@ -0,0 +1,267 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests model-independent block-diffusion rollout transitions.""" + +from absl.testing import absltest +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.diffusion import denoise + + +def _target_logits(targets, vocab_size=32, high=12.0): + return jax.nn.one_hot(targets, vocab_size, dtype=jnp.float32) * high + + +class DiffusionDenoiseTest(absltest.TestCase): + + def test_same_position_generates_partial_blocks_and_preserves_prompt(self): + initial = jnp.asarray([[7, 8, 1, 1, 1, 1, 1, 0]], dtype=jnp.int32) + expected = jnp.asarray([[7, 8, 12, 13, 14, 15, 16, 0]], dtype=jnp.int32) + positions = jnp.arange(8, dtype=jnp.int32)[None, :] + validity = initial != 0 + completion = jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 0]], dtype=jnp.bool_) + + generated = denoise.low_confidence_generate( + lambda _: _target_logits(expected), + initial, + positions, + validity, + completion, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + ) + + np.testing.assert_array_equal(generated, expected) + + def test_shifted_seed_generates_anchors_before_each_block(self): + initial = jnp.asarray([[5, 6, 1, 1, 1, 1, 1, 0]], dtype=jnp.int32) + expected = jnp.asarray([[5, 6, 9, 10, 11, 12, 13, 0]], dtype=jnp.int32) + positions = jnp.arange(8, dtype=jnp.int32)[None, :] + validity = initial != 0 + completion = jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 0]], dtype=jnp.bool_) + + generated = denoise.low_confidence_generate( + lambda _: _target_logits(expected), + initial, + positions, + validity, + completion, + block_size=4, + mask_id=31, + logit_alignment="shifted", + canvas_policy="seed_and_mask", + ) + + np.testing.assert_array_equal(generated, expected) + self.assertEqual(int(generated[0, 4]), 11) + + def test_forced_argmax_progress_completes_low_confidence_rows(self): + initial = jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + + generated = denoise.low_confidence_generate( + lambda canvas: jnp.zeros((*canvas.shape, 8), dtype=jnp.float32), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + ) + + self.assertFalse(bool(jnp.any(generated == 7))) + np.testing.assert_array_equal(generated[:, :1], initial[:, :1]) + + def test_mask_token_is_never_committed(self): + initial = jnp.asarray([[6, 1, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + + def mask_favoring_logits(canvas): + logits = jnp.zeros((*canvas.shape, 8), dtype=jnp.float32) + return logits.at[..., 7].set(100.0).at[..., 5].set(10.0) + + generated = denoise.low_confidence_generate( + mask_favoring_logits, + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + ) + + np.testing.assert_array_equal(generated, [[6, 5, 5, 5]]) + + def test_mask_token_is_excluded_from_shifted_anchor_proposals(self): + initial = jnp.asarray([[2, 3, 4, 5, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(6, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 0, 0, 0, 1, 1]], dtype=jnp.bool_) + + def mask_favoring_logits(canvas): + logits = jnp.zeros((*canvas.shape, 8), dtype=jnp.float32) + return logits.at[..., 7].set(100.0).at[..., 6].set(10.0) + + generated = denoise.low_confidence_generate( + mask_favoring_logits, + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="shifted", + canvas_policy="seed_and_mask", + ) + + np.testing.assert_array_equal(generated, [[2, 3, 4, 5, 6, 6]]) + + def test_rollout_is_jittable(self): + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + expected = jnp.asarray([[4, 5, 6, 7]], dtype=jnp.int32) + + generate = jax.jit( + lambda initial: denoise.low_confidence_generate( + lambda _: _target_logits(expected), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + ) + ) + + np.testing.assert_array_equal(generate(jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32)), expected) + + def test_logical_positions_drive_blocks_after_reordering(self): + positions = jnp.asarray([[0, 4, 1, 5, 2, 6, 3, 7]], dtype=jnp.int32) + expected = positions + 10 + initial = jnp.where(positions < 2, expected, 1) + completion = positions >= 2 + + generated = denoise.low_confidence_generate( + lambda _: _target_logits(expected), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + ) + + np.testing.assert_array_equal(generated, expected) + + def test_forced_ties_are_invariant_to_physical_reordering(self): + completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + + def generate(physical_order): + positions = jnp.asarray(physical_order, dtype=jnp.int32)[None, :] + initial = jnp.take_along_axis(jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32), positions, axis=1) + physical_completion = jnp.take_along_axis(completion, positions, axis=1) + + def low_confidence_logits(canvas): + committed_count = jnp.sum((canvas != 31) & physical_completion, axis=1) + proposed_token = jnp.broadcast_to((10 + committed_count)[:, None], canvas.shape) + return _target_logits(proposed_token, vocab_size=32, high=0.1) + + output = denoise.low_confidence_generate( + low_confidence_logits, + initial, + positions, + jnp.ones_like(physical_completion), + physical_completion, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + ) + inverse_order = jnp.argsort(positions, axis=1) + return jnp.take_along_axis(output, inverse_order, axis=1) + + np.testing.assert_array_equal(generate([0, 1, 2, 3]), generate([0, 3, 1, 2])) + + def test_rejects_non_suffix_and_shifted_origin_completion(self): + positions = np.arange(4, dtype=np.int32)[None, :] + validity = np.ones((1, 4), dtype=bool) + + with self.assertRaisesRegex(ValueError, "contiguous suffix"): + denoise.validate_completion_suffix(positions, validity, np.asarray([[0, 1, 0, 1]], dtype=bool)) + with self.assertRaisesRegex(ValueError, "position zero"): + denoise.validate_completion_suffix( + positions, + validity, + np.asarray([[1, 1, 1, 1]], dtype=bool), + shifted_seed=True, + ) + + def test_rejects_missing_duplicate_or_out_of_range_positions(self): + validity = np.ones((1, 3), dtype=bool) + completion = np.asarray([[0, 1, 1]], dtype=bool) + + for positions in ( + np.asarray([[1, 2, 3]], dtype=np.int32), + np.asarray([[0, 1, 1]], dtype=np.int32), + np.asarray([[0, 1, 5]], dtype=np.int32), + ): + with self.subTest(positions=positions), self.assertRaisesRegex(ValueError, "logical positions"): + denoise.validate_completion_suffix(positions, validity, completion) + + def test_rejects_unsupported_contract_and_short_step_cap(self): + tokens = jnp.ones((1, 4), dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + mask = jnp.ones_like(tokens, dtype=jnp.bool_) + + with self.assertRaisesRegex(ValueError, "supports only"): + denoise.low_confidence_generate( + lambda _: _target_logits(tokens), + tokens, + positions, + mask, + mask, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="seed_and_mask", + ) + with self.assertRaisesRegex(ValueError, "at least block_size"): + denoise.low_confidence_generate( + lambda _: _target_logits(tokens), + tokens, + positions, + mask, + mask, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + max_denoise_steps=3, + ) From b22bc159f0ba3aa447b5b8029f9c650dcb7a1763 Mon Sep 17 00:00:00 2001 From: ethannnnnn Date: Thu, 23 Jul 2026 01:00:33 -0700 Subject: [PATCH 4/6] [maxtext] Add model-aware block-diffusion OPD Add a default-off student-rollout distillation source that prepares fresh block-diffusion rollouts in MaxText and delegates the weighted prepared loss to Tunix. Keep completion, corruption, validity, and loss ownership explicit; score clean generated tokens with a causal teacher; and exclude positions after model-specific stop tokens. Canonicalize multi-turn examples around the final assistant span. Earlier turns become prompt context, while any later conversation turns are truncated from tokens and both input/target segmentation before rollout. Rows without a completion span or prompt still fail closed, so the teacher and student cannot observe future context. Make resume fail closed. Persist the tokenizer, model, objective, optimizer, data-stream, topology, and stop-token contract; replay deterministic HF input before global device placement; reject early exhaustion and incompatible standard-distillation checkpoints. Honor each decoder family's native sharding mode and limit OPD to one inflight computation to control dense-logit memory. The existing dataset-driven distillation path remains the default and lazily avoids the new Tunix diffusion APIs. Test plan: - 209 passed, 14 skipped, 3 deselected, 69 subtests in the focused MaxText suite - 26 passed, 25 skipped, 22 subtests against the paired Tunix OPD head - 4 standard checkpoint restore tests passed through unittest - 45 focused MaxText OPD/input-pipeline tests and 21 subtests passed - Pyink 24.10.1, Ruff, Pylint 10/10, and git diff --check --- src/maxtext/configs/base.yml | 6 + src/maxtext/configs/types.py | 29 + .../input_pipeline/hf_data_processing.py | 1 + .../input_pipeline/input_pipeline_utils.py | 46 ++ .../post_train/distillation/diffusion_opd.py | 622 ++++++++++++++++++ .../distillation/distillation_utils.py | 9 + .../post_train/distillation/train_distill.py | 95 ++- .../unit/maxtext_diffusion_opd_test.py | 539 +++++++++++++++ .../post_training/unit/train_distill_test.py | 40 +- tests/unit/configs_value_test.py | 3 + tests/unit/input_pipeline_utils_test.py | 71 ++ 11 files changed, 1440 insertions(+), 21 deletions(-) create mode 100644 src/maxtext/trainers/post_train/distillation/diffusion_opd.py create mode 100644 tests/post_training/unit/maxtext_diffusion_opd_test.py diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 1091ff7d8d..bd77734d16 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -1321,6 +1321,12 @@ engram_kernel_size: 4 engram_seed: 0 ##### Distillation parameters +distill_data_source: "dataset" # dataset or student_rollout for diffusion on-policy distillation +distill_rollout_algorithm: "low_confidence" +distill_rollout_confidence_threshold: 0.9 +distill_rollout_temperature: 1.0 +distill_rollout_max_denoise_steps: -1 # -1 uses block_diffusion_block_size +distill_rollout_stop_token_ids: [] # Empty uses the tokenizer EOS; set model-specific EOT IDs explicitly distill_alpha: 0.5 distill_temperature: 1.0 # distill_beta is used for cosine similarity loss between intermediate activataitions of out_proj in teacher/student models. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index ccc817b090..4fcdc5f289 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1511,6 +1511,35 @@ class Distillation(BaseModel): description="GCS or local path to the pre-generated ArrayRecord teacher data.", ) + distill_data_source: Literal["dataset", "student_rollout"] = Field( + "dataset", + description="Whether distillation consumes dataset targets or fresh rollouts from the current student.", + ) + distill_rollout_algorithm: Literal["low_confidence"] = Field( + "low_confidence", + description="Block-diffusion rollout algorithm used for on-policy distillation.", + ) + distill_rollout_confidence_threshold: float = Field( + 0.9, + ge=0.0, + le=1.0, + description="Confidence threshold for parallel token commits during diffusion rollout.", + ) + distill_rollout_temperature: float = Field( + 1.0, + gt=0.0, + description="Temperature used to compute rollout token confidence.", + ) + distill_rollout_max_denoise_steps: int = Field( + -1, + ge=-1, + description="Maximum denoising iterations per block; -1 uses block_diffusion_block_size.", + ) + distill_rollout_stop_token_ids: list[int] = Field( + default_factory=list, + description="Generated token IDs that terminate the OPD completion; empty uses the tokenizer EOS ID.", + ) + # --- Loss Params --- distill_alpha: float = Field(0.5, description="Weight for the distillation loss component.") distill_temperature: float = Field(1.0, description="Temperature for distillation softening.") diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py index a047baadbb..9b388778cd 100644 --- a/src/maxtext/input_pipeline/hf_data_processing.py +++ b/src/maxtext/input_pipeline/hf_data_processing.py @@ -68,6 +68,7 @@ def _get_training_objective_transform( completion_only=bool(use_sft and completion_only), seed_first_token=config.block_diffusion_canvas_policy == "seed_and_mask", include_seed_in_loss=config.block_diffusion_logit_alignment == "shifted", + select_last_completion_suffix=getattr(config, "distill_data_source", "dataset") == "student_rollout", ) if objective != "causal_lm": raise ValueError(f"Unsupported training objective: {objective}") diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index 6b294a6dcb..9eead4ac91 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -1005,6 +1005,35 @@ def map(self, element): return shift_and_refine(element, ignored_ids=self.ignored_ids, axis=self.axis) +def _select_last_completion_suffix(validity_mask, completion_mask, axis): + """Selects the final assistant turn and truncates later conversation turns.""" + sequence_length = validity_mask.shape[axis] + moved_validity = np.moveaxis(validity_mask, axis, -1) + moved_completion = np.moveaxis(completion_mask, axis, -1) + flat_validity = moved_validity.reshape(-1, sequence_length) + flat_completion = moved_completion.reshape(-1, sequence_length) + positions = np.arange(sequence_length)[None, :] + + has_completion = flat_completion.any(axis=-1) + last_completion = np.max(np.where(flat_completion, positions, -1), axis=-1) + last_prefix_token = np.max( + np.where((positions < last_completion[:, None]) & ~flat_completion, positions, -1), + axis=-1, + ) + suffix_start = last_prefix_token + 1 + has_prompt = np.any(flat_validity & (positions < suffix_start[:, None]), axis=-1) + nonempty_rows = flat_validity.any(axis=-1) + invalid_rows = nonempty_rows & (~has_completion | ~has_prompt) + if np.any(invalid_rows): + raise ValueError("diffusion OPD requires a completion span after a prompt") + + retained = flat_validity & (positions <= last_completion[:, None]) & has_completion[:, None] + selected = retained & (positions >= suffix_start[:, None]) + selected = selected.reshape(moved_completion.shape) + retained = retained.reshape(moved_validity.shape) + return np.moveaxis(selected, -1, axis), np.moveaxis(retained, -1, axis) + + @dataclasses.dataclass class BlockDiffusionCorruption(grain.RandomMapTransform): """Builds explicit completion, corruption, and loss masks for block diffusion.""" @@ -1018,6 +1047,7 @@ def __init__( completion_only: bool = False, seed_first_token: bool = False, include_seed_in_loss: bool = False, + select_last_completion_suffix: bool = False, ): if block_size <= 0: raise ValueError(f"block_size must be positive, got {block_size}") @@ -1025,6 +1055,8 @@ def __init__( raise ValueError(f"min_noise must satisfy 0 < min_noise <= 1, got {min_noise}") if include_seed_in_loss and not seed_first_token: raise ValueError("include_seed_in_loss requires seed_first_token=True") + if select_last_completion_suffix and not completion_only: + raise ValueError("select_last_completion_suffix requires completion_only=True") self.block_size = block_size self.mask_id = mask_id self.min_noise = min_noise @@ -1032,6 +1064,7 @@ def __init__( self.completion_only = completion_only self.seed_first_token = seed_first_token self.include_seed_in_loss = include_seed_in_loss + self.select_last_completion_suffix = select_last_completion_suffix def random_map(self, element, rng: np.random.Generator): inputs = np.asarray(element["inputs"]) @@ -1053,6 +1086,13 @@ def random_map(self, element, rng: np.random.Generator): if completion_mask.shape != inputs.shape: raise ValueError(f"completion_mask must match inputs shape; received {completion_mask.shape} and {inputs.shape}") completion_mask &= validity_mask + if self.select_last_completion_suffix: + completion_mask, validity_mask = _select_last_completion_suffix(validity_mask, completion_mask, axis) + inputs = np.where(validity_mask, inputs, inputs.dtype.type(0)) + targets = np.where(validity_mask, targets, targets.dtype.type(0)) + targets_segmentation = np.where( + validity_mask, targets_segmentation, targets_segmentation.dtype.type(0) + ) eligible_mask = completion_mask if self.completion_only else validity_mask moved_eligible = np.moveaxis(eligible_mask, axis, -1) @@ -1083,6 +1123,12 @@ def random_map(self, element, rng: np.random.Generator): output = dict(element) output["inputs"] = np.where(corruption_mask, inputs.dtype.type(self.mask_id), inputs) output["targets"] = targets + if self.select_last_completion_suffix: + inputs_segmentation = np.asarray(element["inputs_segmentation"]) + output["inputs_segmentation"] = np.where( + validity_mask, inputs_segmentation, inputs_segmentation.dtype.type(0) + ) + output["targets_segmentation"] = targets_segmentation output["completion_mask"] = completion_mask.astype(targets_segmentation.dtype) output["corruption_mask"] = corruption_mask.astype(targets_segmentation.dtype) output["targets_loss_mask"] = targets_loss_mask.astype(targets_segmentation.dtype) diff --git a/src/maxtext/trainers/post_train/distillation/diffusion_opd.py b/src/maxtext/trainers/post_train/distillation/diffusion_opd.py new file mode 100644 index 0000000000..1a858ce9b5 --- /dev/null +++ b/src/maxtext/trainers/post_train/distillation/diffusion_opd.py @@ -0,0 +1,622 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MaxText model-aware preparation for Tunix diffusion OPD.""" + +import math + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.diffusion import denoise +from maxtext.diffusion import scoring +from maxtext.trainers.post_train.distillation import distillation_utils +from maxtext.utils import max_logging +from tunix.diffusion import types as diffusion_types +from tunix.distillation import diffusion as diffusion_distillation +from tunix.distillation import diffusion_opd as tunix_diffusion_opd +from tunix.sft import peft_trainer + + +def _resolve_logit_alignment(config) -> str: + if getattr(config, "training_objective", "causal_lm") == "block_diffusion": + return config.block_diffusion_logit_alignment + return "shifted" + + +def create_target_aligned_logits_fn(config, *, enable_dropout: bool): + """Builds a MaxText forward adapter satisfying Tunix's logits contract.""" + alignment = _resolve_logit_alignment(config) + + def logits_fn(model, model_inputs): + logits = model( + decoder_input_tokens=model_inputs["input_tokens"], + decoder_positions=model_inputs["positions"], + decoder_segment_ids=model_inputs["decoder_segment_ids"], + enable_dropout=enable_dropout, + decoder_target_tokens=model_inputs["targets"], + decoder_target_mask=model_inputs["targets_segmentation"], + ) + return scoring.align_logits_to_targets( + logits, + alignment, + model_inputs["positions"], + model_inputs["targets_segmentation"] != 0, + ) + + return logits_fn + + +def create_rollout_fn(config): + """Builds one compiled student rollout with deterministic dropout settings.""" + logits_fn = create_target_aligned_logits_fn(config, enable_dropout=False) + max_denoise_steps = config.distill_rollout_max_denoise_steps + if max_denoise_steps == -1: + max_denoise_steps = config.block_diffusion_block_size + + @nnx.jit + def rollout(model, initial_tokens, positions, decoder_segment_ids, completion_mask): + graphdef, params, rest = nnx.split(model, nnx.Param, ...) + validity_mask = decoder_segment_ids != 0 + + def canvas_logits(canvas): + # The denoising state machine uses nested JAX control flow. Reconstructing + # a local model at that trace level avoids mutating outer NNX Variables; + # rollout uses dropout=False, so its transient non-parameter state is + # intentionally discarded after each proposal forward. + local_model = nnx.merge(graphdef, params, rest, copy=True) + return logits_fn( + local_model, + { + "input_tokens": canvas, + "positions": positions, + "decoder_segment_ids": decoder_segment_ids, + "targets": canvas, + "targets_segmentation": decoder_segment_ids, + }, + ) + + return denoise.low_confidence_generate( + canvas_logits, + initial_tokens, + positions, + validity_mask, + completion_mask, + block_size=config.block_diffusion_block_size, + mask_id=config.block_diffusion_mask_id, + logit_alignment=config.block_diffusion_logit_alignment, + canvas_policy=config.block_diffusion_canvas_policy, + confidence_threshold=config.distill_rollout_confidence_threshold, + temperature=config.distill_rollout_temperature, + max_denoise_steps=max_denoise_steps, + ) + + return rollout + + +def create_teacher_score_fn(config): + """Builds a compiled frozen-teacher scoring function.""" + logits_fn = create_target_aligned_logits_fn(config, enable_dropout=False) + + @nnx.jit + def score(model, model_inputs): + graphdef, params, rest = nnx.split(model, nnx.Param, ...) + local_model = nnx.merge(graphdef, params, rest, copy=True) + return logits_fn(local_model, model_inputs) + + return score + + +def _concrete_numpy(value): + if isinstance(value, jax.core.Tracer): + return None + if isinstance(value, jax.Array) and not value.is_fully_addressable: + return None + return np.asarray(value) + + +def _validate_input_masks( + positions, + validity_mask, + completion_mask, + corruption_mask, + loss_mask, + *, + logit_alignment, + block_size, +): + """Validates role and loss ownership before launching model computation.""" + shapes = { + "positions": tuple(positions.shape), + "validity_mask": tuple(validity_mask.shape), + "completion_mask": tuple(completion_mask.shape), + "corruption_mask": tuple(corruption_mask.shape), + "targets_loss_mask": tuple(loss_mask.shape), + } + if len(set(shapes.values())) != 1: + raise ValueError(f"diffusion OPD masks must have identical shapes; received {shapes}") + concrete = [_concrete_numpy(value) for value in (positions, validity_mask, completion_mask, corruption_mask, loss_mask)] + if any(value is None for value in concrete): + return + concrete_positions = np.asarray(concrete[0]) + validity, completion, corruption, weights = (np.asarray(value, dtype=bool) for value in concrete[1:]) + if np.any(completion & ~validity): + raise ValueError("completion_mask must be a subset of valid target positions") + if np.any(corruption & ~completion): + raise ValueError("corruption_mask must be a subset of completion_mask") + if np.any(weights & ~completion): + raise ValueError("targets_loss_mask must be a subset of completion_mask") + if np.any(corruption & ~weights): + raise ValueError("every corrupted position must carry a positive target loss weight") + safe_clean_targets = np.zeros_like(weights) + if logit_alignment == "shifted": + safe_clean_targets = (concrete_positions > 0) & (concrete_positions % block_size == 0) + if np.any(weights & ~corruption & ~safe_clean_targets): + raise ValueError( + "weighted clean student targets are allowed only at shifted logical block anchors; " + "same-position targets must be corrupted" + ) + if np.any(weights & (concrete_positions == 0)): + raise ValueError("causal teacher distillation must not weight logical position zero") + + +def _safe_loss_weights( + positions, validity_mask, completion_mask, corruption_mask, loss_weights, *, alignment, block_size +): + """Masks unsafe rows and target positions using device-side invariants.""" + sequence_length = positions.shape[1] + batch_size = positions.shape[0] + validity_mask = jnp.asarray(validity_mask, dtype=jnp.bool_) + completion_mask = jnp.asarray(completion_mask, dtype=jnp.bool_) + corruption_mask = jnp.asarray(corruption_mask, dtype=jnp.bool_) + loss_weights = jnp.asarray(loss_weights) + positions = jnp.asarray(positions, dtype=jnp.int32) + weighted_positions = loss_weights > 0 + completion_is_valid = jnp.all(~completion_mask | validity_mask, axis=1) + corruption_is_completion = jnp.all(~corruption_mask | completion_mask, axis=1) + loss_is_completion = jnp.all(~weighted_positions | completion_mask, axis=1) + corruption_is_weighted = jnp.all(~corruption_mask | weighted_positions, axis=1) + positions_in_range = jnp.all(~validity_mask | ((positions >= 0) & (positions < sequence_length)), axis=1) + valid_completion_mask = completion_mask & validity_mask + valid_corruption_mask = corruption_mask & valid_completion_mask + safe_positions = jnp.clip(positions, 0, sequence_length - 1) + row_indices = jnp.broadcast_to(jnp.arange(batch_size, dtype=jnp.int32)[:, None], positions.shape) + + logical_counts = jnp.zeros_like(positions).at[row_indices, safe_positions].add(validity_mask.astype(jnp.int32)) + valid_count = jnp.sum(validity_mask, axis=1) + expected_counts = jnp.arange(sequence_length, dtype=jnp.int32)[None, :] < valid_count[:, None] + positions_are_canonical = jnp.all(logical_counts == expected_counts.astype(jnp.int32), axis=1) + + logical_completion = jnp.zeros_like(valid_completion_mask).at[row_indices, safe_positions].max(valid_completion_mask) + seen_completion = jax.lax.associative_scan(jnp.logical_or, logical_completion, axis=1) + invalid_future_context = seen_completion & expected_counts & ~logical_completion + has_prompt_origin = expected_counts[:, 0] & ~logical_completion[:, 0] + has_completion = jnp.any(logical_completion, axis=1) + row_is_safe = ( + positions_in_range + & positions_are_canonical + & completion_is_valid + & corruption_is_completion + & loss_is_completion + & corruption_is_weighted + & ~jnp.any(invalid_future_context, axis=1) + & has_prompt_origin + & has_completion + ) + numeric_weights_are_safe = jnp.all(jnp.isfinite(loss_weights) & (loss_weights >= 0), axis=1) + row_is_safe &= numeric_weights_are_safe + + allowed_positions = valid_completion_mask & (positions != 0) + if alignment == "same_position": + allowed_positions &= valid_corruption_mask + else: + shifted_anchors = (positions > 0) & (positions % block_size == 0) + allowed_positions &= valid_corruption_mask | shifted_anchors + row_is_safe &= jnp.all(~weighted_positions | allowed_positions, axis=1) + return jnp.where(row_is_safe[:, None] & allowed_positions, loss_weights, 0.0) + + +def _completion_through_first_stop(generated_tokens, completion_mask, stop_token_ids): + """Keeps completion positions through the first generated stop token, inclusive.""" + if not stop_token_ids: + return completion_mask + stop_positions = completion_mask & jnp.isin(generated_tokens, jnp.asarray(stop_token_ids, generated_tokens.dtype)) + seen_stop = jax.lax.associative_scan(jnp.logical_or, stop_positions, axis=1) + seen_stop_before = jnp.concatenate([jnp.zeros_like(seen_stop[:, :1]), seen_stop[:, :-1]], axis=1) + return completion_mask & ~seen_stop_before + + +def prepare_diffusion_opd_batch( + input_data: distillation_utils.MaxTextTrainingInput, + student_model: nnx.Module, + teacher_model: nnx.Module, + *, + rollout_fn, + teacher_score_fn, + mask_id: int, + stop_token_ids: tuple[int, ...] = (), + shifted_seed: bool = False, + logit_alignment: str = "same_position", + block_size: int = 1, +) -> diffusion_distillation.DiffusionDistillationBatch: + """Creates a fresh rollout and teacher scores for one OPD step.""" + required_fields = { + "targets": input_data.targets, + "positions": input_data.positions, + "completion_mask": input_data.completion_mask, + "corruption_mask": input_data.corruption_mask, + "targets_loss_mask": input_data.targets_loss_mask, + } + missing = sorted(name for name, value in required_fields.items() if value is None) + if missing: + raise ValueError(f"diffusion OPD requires explicit batch fields; missing {missing}") + targets = input_data.targets + positions = input_data.positions + decoder_segment_ids = input_data.decoder_segment_ids + if decoder_segment_ids is None: + decoder_segment_ids = jnp.asarray(input_data.input_mask, dtype=jnp.int32) + validity_mask = decoder_segment_ids != 0 + completion_mask = jnp.asarray(input_data.completion_mask, dtype=jnp.bool_) + corruption_mask = jnp.asarray(input_data.corruption_mask, dtype=jnp.bool_) + loss_weights = jnp.asarray(input_data.targets_loss_mask, dtype=jnp.float32) + _validate_input_masks( + positions, + validity_mask, + completion_mask, + corruption_mask, + loss_weights, + logit_alignment=logit_alignment, + block_size=block_size, + ) + completion_mask &= validity_mask + corruption_mask &= completion_mask + denoise.validate_completion_suffix( + positions, + validity_mask, + completion_mask, + shifted_seed=shifted_seed, + ) + + generated_tokens = rollout_fn( + student_model, + targets, + positions, + decoder_segment_ids, + completion_mask, + ) + completion_mask = _completion_through_first_stop(generated_tokens, completion_mask, stop_token_ids) + validity_mask &= ~jnp.asarray(input_data.completion_mask, dtype=jnp.bool_) | completion_mask + decoder_segment_ids = jnp.where(validity_mask, decoder_segment_ids, 0) + corruption_mask &= completion_mask + loss_weights = jnp.where(completion_mask, loss_weights, 0.0) + student_inputs = jnp.where(corruption_mask, jnp.asarray(mask_id, generated_tokens.dtype), generated_tokens) + loss_weights = _safe_loss_weights( + positions, + validity_mask, + completion_mask, + corruption_mask, + loss_weights, + alignment=logit_alignment, + block_size=block_size, + ) + model_inputs = { + "input_tokens": student_inputs, + "positions": positions, + "decoder_segment_ids": decoder_segment_ids, + "targets": generated_tokens, + "targets_segmentation": decoder_segment_ids, + } + student_batch = diffusion_types.DiffusionTokenBatch.create( + model_inputs=model_inputs, + target_ids=generated_tokens, + loss_weights=loss_weights, + ) + teacher_inputs = dict(model_inputs) + teacher_inputs["input_tokens"] = generated_tokens + teacher_logits = teacher_score_fn(teacher_model, teacher_inputs) + return diffusion_distillation.DiffusionDistillationBatch.create( + student_batch=student_batch, + teacher_logits=teacher_logits, + ) + + +def validate_diffusion_opd_configs(student_config, teacher_config, *, is_offline: bool) -> bool: + """Validates OPD-only constraints and returns whether OPD is enabled.""" + if student_config.distill_data_source != "student_rollout": + return False + if is_offline: + raise ValueError("diffusion OPD requires an online teacher; offline_data_dir is not supported") + if student_config.training_objective != "block_diffusion": + raise ValueError("diffusion OPD requires student training_objective='block_diffusion'") + if not student_config.use_sft or not student_config.sft_train_on_completion_only: + raise ValueError("diffusion OPD requires completion-only SFT data with an explicit prompt/completion boundary") + if teacher_config.training_objective != "causal_lm" or teacher_config.attention_type in { + "block_diffusion", + "full", + "compressed", + }: + raise ValueError("diffusion OPD currently requires a causal teacher to score clean student rollouts") + if getattr(teacher_config, "mtp_num_layers", 0) > 0: + raise ValueError("diffusion OPD does not support an MTP teacher") + if student_config.vocab_size != teacher_config.vocab_size: + raise ValueError("diffusion OPD requires matching student and teacher vocabularies") + student_tokenizer = ( + str(getattr(student_config, "tokenizer_path", "")), + getattr(getattr(student_config, "tokenizer_type", ""), "value", getattr(student_config, "tokenizer_type", "")), + ) + teacher_tokenizer = ( + str(getattr(teacher_config, "tokenizer_path", "")), + getattr(getattr(teacher_config, "tokenizer_type", ""), "value", getattr(teacher_config, "tokenizer_type", "")), + ) + if student_tokenizer != teacher_tokenizer: + raise ValueError( + "diffusion OPD requires identical student and teacher tokenizer_path/tokenizer_type; " + "token ID remapping is not supported" + ) + if student_tokenizer[1] != "huggingface": + raise ValueError("diffusion OPD requires tokenizer_type='huggingface' to match the HF input pipeline") + if student_config.distill_beta != 0.0 or student_config.distill_beta_end is not None: + raise ValueError("diffusion OPD does not support feature distillation") + if student_config.distill_alpha_schedule != "constant" or student_config.distill_temperature_schedule != "constant": + raise ValueError("diffusion OPD currently requires constant alpha and temperature") + if student_config.distill_alpha_end is not None or student_config.distill_temperature_end is not None: + raise ValueError("diffusion OPD does not accept scheduled alpha or temperature end values") + if not math.isfinite(student_config.distill_alpha) or not 0.0 <= student_config.distill_alpha <= 1.0: + raise ValueError("distill_alpha must be finite and in [0, 1] for diffusion OPD") + if not math.isfinite(student_config.distill_temperature) or student_config.distill_temperature <= 0.0: + raise ValueError("distill_temperature must be finite and positive for diffusion OPD") + if not math.isfinite(student_config.distill_rollout_temperature): + raise ValueError("distill_rollout_temperature must be finite for diffusion OPD") + if not math.isfinite(student_config.distill_rollout_confidence_threshold): + raise ValueError("distill_rollout_confidence_threshold must be finite for diffusion OPD") + if student_config.eval_interval > 0: + raise ValueError("diffusion OPD evaluation is not supported yet; set eval_interval=-1") + if student_config.steps <= 0: + raise ValueError("diffusion OPD requires a positive finite number of training steps") + if student_config.gradient_accumulation_steps > 1 and not student_config.use_tunix_gradient_accumulation: + raise ValueError( + "diffusion OPD with gradient accumulation requires use_tunix_gradient_accumulation=True " + "so the HF batch is divided into Tunix microbatches exactly once" + ) + if getattr(student_config, "generate_padding_batch_train", False): + raise ValueError("diffusion OPD does not support generate_padding_batch_train") + if getattr(student_config, "elastic_enabled", False): + raise ValueError("diffusion OPD does not yet support elastic data loading or topology changes") + if getattr(student_config, "expansion_factor_real_data", -1.0) != -1.0: + raise ValueError("diffusion OPD does not support expansion_factor_real_data") + if getattr(student_config, "enable_rampup_batch_size", False): + raise ValueError("diffusion OPD does not support ramp-up batch sizing") + if student_config.learn_to_init_mode: + raise ValueError("diffusion OPD is not compatible with learn_to_init_mode") + if getattr(student_config, "student_params_to_update", None): + raise ValueError("diffusion OPD does not yet support student_params_to_update filtering") + max_steps = student_config.distill_rollout_max_denoise_steps + if max_steps != -1 and max_steps < student_config.block_diffusion_block_size: + raise ValueError("distill_rollout_max_denoise_steps must be -1 or at least block_diffusion_block_size") + stop_token_ids = student_config.distill_rollout_stop_token_ids + if len(set(stop_token_ids)) != len(stop_token_ids) or any( + token_id < 0 or token_id >= student_config.vocab_size for token_id in stop_token_ids + ): + raise ValueError("distill_rollout_stop_token_ids must contain unique IDs in the student vocabulary") + return True + + +def _checkpoint_contract(student_config, teacher_config, stop_token_ids=()): + """Builds the semantic identity that must remain stable across resume.""" + return { + "version": 1, + "student_model": str(getattr(student_config, "model_name", "")), + "student_tokenizer": str(getattr(student_config, "tokenizer_path", "")), + "student_tokenizer_type": str( + getattr(getattr(student_config, "tokenizer_type", ""), "value", getattr(student_config, "tokenizer_type", "")) + ), + "teacher_model": str(getattr(teacher_config, "model_name", "")), + "teacher_checkpoint": str(getattr(teacher_config, "load_parameters_path", "")), + "teacher_tokenizer": str(getattr(teacher_config, "tokenizer_path", "")), + "teacher_tokenizer_type": str( + getattr(getattr(teacher_config, "tokenizer_type", ""), "value", getattr(teacher_config, "tokenizer_type", "")) + ), + "vocab_size": int(student_config.vocab_size), + "block_size": int(student_config.block_diffusion_block_size), + "mask_id": int(student_config.block_diffusion_mask_id), + "stop_token_ids": [int(token_id) for token_id in stop_token_ids], + "min_noise": float(student_config.block_diffusion_min_noise), + "logit_alignment": student_config.block_diffusion_logit_alignment, + "canvas_policy": student_config.block_diffusion_canvas_policy, + "rollout_algorithm": student_config.distill_rollout_algorithm, + "rollout_confidence_threshold": float(student_config.distill_rollout_confidence_threshold), + "rollout_temperature": float(student_config.distill_rollout_temperature), + "rollout_max_denoise_steps": int(student_config.distill_rollout_max_denoise_steps), + "enable_dropout": bool(student_config.enable_dropout), + "dropout_rate": float(getattr(student_config, "dropout_rate", 0.0)), + "steps": int(student_config.steps), + "optimizer_type": str(student_config.opt_type), + "learning_rate": float(student_config.learning_rate), + "warmup_steps_fraction": float(student_config.warmup_steps_fraction), + "learning_rate_final_fraction": float(student_config.learning_rate_final_fraction), + "gradient_clipping_threshold": float(student_config.gradient_clipping_threshold), + "adam_b1": float(getattr(student_config, "adam_b1", 0.0)), + "adam_b2": float(getattr(student_config, "adam_b2", 0.0)), + "adam_eps": float(getattr(student_config, "adam_eps", 0.0)), + "adam_eps_root": float(getattr(student_config, "adam_eps_root", 0.0)), + "adam_weight_decay": float(getattr(student_config, "adam_weight_decay", 0.0)), + "adamw_mask": str(getattr(student_config, "adamw_mask", "")), + "mu_dtype": str(getattr(student_config, "mu_dtype", "")), + "skip_step_on_spikes": bool(getattr(student_config, "skip_step_on_spikes", False)), + "skip_step_interval": int(getattr(student_config, "skip_step_interval", 0)), + "skip_step_scaling_factor": float(getattr(student_config, "skip_step_scaling_factor", 0.0)), + "trainable_parameters_mask": str(getattr(student_config, "trainable_parameters_mask", "")), + "distill_temperature": float(student_config.distill_temperature), + "soft_loss_weight": float(student_config.distill_alpha), + "hard_loss_weight": float(1.0 - student_config.distill_alpha), + "dataset_type": str(student_config.dataset_type), + "hf_path": str(getattr(student_config, "hf_path", "")), + "hf_name": str(getattr(student_config, "hf_name", "")), + "hf_data_dir": str(getattr(student_config, "hf_data_dir", "")), + "hf_train_files": str(getattr(student_config, "hf_train_files", "")), + "train_split": str(getattr(student_config, "train_split", "")), + "train_data_columns": str(getattr(student_config, "train_data_columns", "")), + "tokenize_train_data": bool(getattr(student_config, "tokenize_train_data", True)), + "enable_data_shuffling": bool(getattr(student_config, "enable_data_shuffling", False)), + "data_shuffle_seed": int(getattr(student_config, "data_shuffle_seed", 0)), + "num_epoch": int(getattr(student_config, "num_epoch", 1)), + "max_target_length": int(student_config.max_target_length), + "add_bos": bool(getattr(student_config, "add_bos", True)), + "add_eos": bool(getattr(student_config, "add_eos", True)), + "packing": bool(getattr(student_config, "packing", False)), + "generate_padding_batch_train": bool(getattr(student_config, "generate_padding_batch_train", False)), + "chat_template_path": str(getattr(student_config, "chat_template_path", "")), + "chat_template": str(getattr(student_config, "chat_template", "")), + "formatting_func_path": str(getattr(student_config, "formatting_func_path", "")), + "formatting_func_kwargs": str(getattr(student_config, "formatting_func_kwargs", "")), + "gradient_accumulation_steps": int(student_config.gradient_accumulation_steps), + "use_tunix_gradient_accumulation": bool(student_config.use_tunix_gradient_accumulation), + "global_batch_size_to_load": int(student_config.global_batch_size_to_load), + "global_batch_size_to_train_on": int(getattr(student_config, "global_batch_size_to_train_on", 0)), + "data_sharding": str(getattr(student_config, "data_sharding", "")), + "expansion_factor_real_data": float(getattr(student_config, "expansion_factor_real_data", -1.0)), + "process_count": int(jax.process_count()), + "local_device_count": int(jax.local_device_count()), + } + + +def replay_and_bound_iterator( + iterator, + *, + iter_steps: int, + train_steps: int, + max_steps: int, + accumulation_steps: int, + replay_iterator=None, +): + """Replays a deterministic HF stream and bounds it to remaining updates. + + `replay_iterator` may be the host-local source underlying `iterator`. Advancing + it avoids materializing and device-putting global arrays for skipped batches. + """ + if iter_steps < 0 or train_steps < 0 or accumulation_steps < 1: + raise ValueError("restored iterator and training step counts must be nonnegative") + if train_steps > max_steps: + raise ValueError(f"restored train_steps ({train_steps}) exceeds configured max_steps ({max_steps})") + if iter_steps != train_steps * accumulation_steps: + raise ValueError("restored iter_steps must equal train_steps * accumulation_steps") + remaining_updates = max_steps - train_steps + if remaining_updates == 0: + return iter(()) + replay_source = iterator if replay_iterator is None else replay_iterator + for replayed_steps in range(iter_steps): + try: + next(replay_source) + except StopIteration as exc: + raise ValueError( + f"diffusion OPD HF replay exhausted after {replayed_steps} of {iter_steps} required microbatches" + ) from exc + remaining_microbatches = remaining_updates * accumulation_steps + + def take_remaining_exactly(): + for consumed_steps in range(remaining_microbatches): + try: + yield next(iterator) + except StopIteration as exc: + raise ValueError( + "diffusion OPD HF stream exhausted after " + f"{consumed_steps} of {remaining_microbatches} required remaining microbatches" + ) from exc + + return take_remaining_exactly() + + +class MaxTextDiffusionOPDTrainer(peft_trainer.PeftTrainer): + """Tunix trainer that prepares fresh MaxText rollouts before every JIT step.""" + + def __init__( + self, + *, + student_model, + teacher_model, + optimizer, + training_config, + student_config, + teacher_config, + eos_id, + ): + super().__init__(model=student_model, optimizer=optimizer, training_config=training_config) + if teacher_model is None: + raise ValueError("diffusion OPD requires a teacher model") + configured_stop_ids = tuple(student_config.distill_rollout_stop_token_ids) + if not configured_stop_ids and (eos_id is None or eos_id < 0 or eos_id >= student_config.vocab_size): + raise ValueError("diffusion OPD requires a valid tokenizer EOS token ID") + self.teacher_model = teacher_model + self.student_config = student_config + self.stop_token_ids = configured_stop_ids or (int(eos_id),) + self.checkpoint_contract = _checkpoint_contract(student_config, teacher_config, self.stop_token_ids) + self.rollout_fn = create_rollout_fn(student_config) + self.teacher_score_fn = create_teacher_score_fn(teacher_config) + student_logits_fn = create_target_aligned_logits_fn(student_config, enable_dropout=student_config.enable_dropout) + tunix_diffusion_opd.configure_prepared_diffusion_opd( + self, + lambda batch: batch, + student_logits_fn, + temperature=student_config.distill_temperature, + soft_loss_weight=student_config.distill_alpha, + hard_loss_weight=1.0 - student_config.distill_alpha, + ) + + def _prepare_inputs(self, input_data): + return prepare_diffusion_opd_batch( + input_data, + self.model, + self.teacher_model, + rollout_fn=self.rollout_fn, + teacher_score_fn=self.teacher_score_fn, + mask_id=self.student_config.block_diffusion_mask_id, + stop_token_ids=self.stop_token_ids, + shifted_seed=self.student_config.block_diffusion_logit_alignment == "shifted", + logit_alignment=self.student_config.block_diffusion_logit_alignment, + block_size=self.student_config.block_diffusion_block_size, + ) + + def setup_checkpoint_manager_and_restore(self, raw_train_iter, config): + """Restores model state; HF stream position is replayed deterministically.""" + if self.checkpoint_manager is not None: + self.checkpoint_manager.close() + self.checkpoint_manager = distillation_utils.MaxTextCheckpointManager( + raw_iterator=None, + root_directory=config.checkpoint_dir, + student_config=config, + options=self.config.checkpointing_options, + ) + self._train_steps, self._restored_custom_metadata = self.checkpoint_manager.maybe_restore( + self.model, + self.optimizer, + restore_only_lora_params=getattr(self, "_lora_enabled", False), + ) + if self._train_steps > 0: + restored_contract = self._restored_custom_metadata.get("diffusion_opd_contract") + if restored_contract is None: + raise ValueError("diffusion OPD checkpoint is missing its semantic training contract") + if restored_contract != self.checkpoint_contract: + raise ValueError( + "diffusion OPD checkpoint contract does not match the current student, teacher, tokenizer, or objective" + ) + self._iter_steps = self._train_steps * self.config.get_with_default("gradient_accumulation_steps", 1) + if self._iter_steps: + max_logging.log(f"Diffusion OPD will replay {self._iter_steps} deterministic HF microbatches after restore.") + return raw_train_iter + + def custom_checkpoint_metadata(self): + """Returns the immutable semantic contract persisted with each checkpoint.""" + return {"diffusion_opd_contract": self.checkpoint_contract} diff --git a/src/maxtext/trainers/post_train/distillation/distillation_utils.py b/src/maxtext/trainers/post_train/distillation/distillation_utils.py index cfa7fdfeb2..e8e9bba025 100644 --- a/src/maxtext/trainers/post_train/distillation/distillation_utils.py +++ b/src/maxtext/trainers/post_train/distillation/distillation_utils.py @@ -69,6 +69,12 @@ class MaxTextTrainingInput(peft_trainer.TrainingInput): targets_position: jax.Array | None = None #: Segment IDs for packed target tokens. targets_segmentation: jax.Array | None = None + #: Role-derived positions eligible for completion generation. + completion_mask: jax.Array | None = None + #: Positions replaced by the diffusion mask token for the student forward pass. + corruption_mask: jax.Array | None = None + #: Per-target weights for target-aligned diffusion objectives. + targets_loss_mask: jax.Array | None = None #: Top-K logits from the teacher model. top_k_logits: jax.Array | None = None top_k_indices: jax.Array | None = None @@ -132,6 +138,9 @@ def __next__(self) -> MaxTextTrainingInput: targets=batch["targets"], targets_position=targets_position, targets_segmentation=targets_segmentation, + completion_mask=batch.get("completion_mask"), + corruption_mask=batch.get("corruption_mask"), + targets_loss_mask=batch.get("targets_loss_mask"), top_k_logits=batch.get("top_k_logits"), top_k_indices=batch.get("top_k_indices"), ) diff --git a/src/maxtext/trainers/post_train/distillation/train_distill.py b/src/maxtext/trainers/post_train/distillation/train_distill.py index f31b43b921..84ae819a3c 100644 --- a/src/maxtext/trainers/post_train/distillation/train_distill.py +++ b/src/maxtext/trainers/post_train/distillation/train_distill.py @@ -422,6 +422,9 @@ def _prepare_inputs( targets=input_data.targets, targets_position=input_data.targets_position, targets_segmentation=input_data.targets_segmentation, + completion_mask=input_data.completion_mask, + corruption_mask=input_data.corruption_mask, + targets_loss_mask=input_data.targets_loss_mask, top_k_logits=input_data.top_k_logits, top_k_indices=input_data.top_k_indices, ) @@ -532,6 +535,11 @@ def setup_checkpoint_manager_and_restore(self, raw_train_iter, config): self.optimizer, restore_only_lora_params=getattr(self, "_lora_enabled", False), ) + if "diffusion_opd_contract" in (self._restored_custom_metadata or {}): + raise ValueError( + "standard distillation cannot resume a diffusion OPD checkpoint; " + "keep distill_data_source='student_rollout' or start from a compatible checkpoint" + ) grad_accum_steps = self.config.get_with_default("gradient_accumulation_steps", 1) self._iter_steps = self._train_steps * grad_accum_steps @@ -574,14 +582,14 @@ def build_training_components( is_offline: bool = False, offline_data_dir: str | None = None, ): - """Builds and returns the strategy, optimizer, and training config objects. + """Builds the strategy, optimizer, training config, and tokenizer EOS ID. Args: student_config: Configuration object for the Student model. teacher_config: Configuration object for the Teacher model. Returns: - A tuple of (DistillationStrategy, Optimizer, TrainingConfig). + A tuple of (DistillationStrategy, Optimizer, TrainingConfig, eos_id). """ # 2. Load Tokenizer Info tok = tokenizer.build_tokenizer( @@ -642,6 +650,9 @@ def build_training_components( log_dir=student_config.tensorboard_dir, flush_every_n_steps=student_config.log_period ) + opd_training_options = ( + {"max_inflight_computations": 1} if student_config.distill_data_source == "student_rollout" else {} + ) train_config = peft_trainer.TrainingConfig( max_steps=student_config.steps, eval_every_n_steps=student_config.eval_interval, @@ -651,9 +662,10 @@ def build_training_components( checkpointing_options=checkpointing_options, gradient_accumulation_steps=student_config.gradient_accumulation_steps, data_sharding_axis=tuple(student_config.data_sharding), + **opd_training_options, ) - return strategy, optimizer, train_config + return strategy, optimizer, train_config, tok.eos_id def train_distill( @@ -677,9 +689,19 @@ def train_distill( f"Vocab size mismatch! Student: {student_config.vocab_size}, Teacher: {teacher_config.vocab_size}. " "Distillation requires matching vocabularies." ) + diffusion_opd_enabled = student_config.distill_data_source == "student_rollout" + diffusion_opd_lib = None + if diffusion_opd_enabled: + # Keep the standard distillation import path compatible with the currently + # pinned Tunix release. New Tunix diffusion APIs are required only when OPD + # is explicitly selected. + # pylint: disable-next=import-outside-toplevel + from maxtext.trainers.post_train.distillation import diffusion_opd as diffusion_opd_lib + + diffusion_opd_lib.validate_diffusion_opd_configs(student_config, teacher_config, is_offline=is_offline) # Build Training Components (No hardware context required) - strategy, optimizer, train_config = build_training_components( + strategy, optimizer, train_config, eos_id = build_training_components( student_config, teacher_config, is_offline, offline_data_dir ) @@ -724,21 +746,33 @@ def student_freeze_param_fn(path) -> bool: ) student_model.train() - model_bundle = ModelBundle(teacher_model, student_model) - - # 3. Initialize Trainer - trainer = MaxTextDistillationTrainer( - model=model_bundle, - strategy=strategy, - optimizer=optimizer, - training_config=train_config, - student_config=student_config, - teacher_config=teacher_config, - is_offline=is_offline, - student_freeze_param_filter=student_freeze_param_fn if student_params_to_update else None, - ) + model_bundle = None + if diffusion_opd_enabled: + assert diffusion_opd_lib is not None + trainer = diffusion_opd_lib.MaxTextDiffusionOPDTrainer( + student_model=student_model, + teacher_model=teacher_model, + optimizer=optimizer, + training_config=train_config, + student_config=student_config, + teacher_config=teacher_config, + eos_id=eos_id, + ) + else: + model_bundle = ModelBundle(teacher_model, student_model) + trainer = MaxTextDistillationTrainer( + model=model_bundle, + strategy=strategy, + optimizer=optimizer, + training_config=train_config, + student_config=student_config, + teacher_config=teacher_config, + is_offline=is_offline, + student_freeze_param_filter=student_freeze_param_fn if student_params_to_update else None, + ) trainer.is_managed_externally = True - trainer._has_aux = True # pylint: disable=protected-access + if not diffusion_opd_enabled: + trainer._has_aux = True # pylint: disable=protected-access if is_offline: max_logging.log("Initializing Data Iterators via MaxText pipeline...") @@ -759,7 +793,10 @@ def student_freeze_param_fn(path) -> bool: # Sync the ModelBundle step counter with the restored training step so that # loss weight schedules resume from the correct position after checkpoint restore. - model_bundle.training_step.set_value(jnp.array(trainer._train_steps, dtype=jnp.int32)) # pylint: disable=protected-access + if model_bundle is not None: + model_bundle.training_step.set_value( # pylint: disable=protected-access + jnp.array(trainer._train_steps, dtype=jnp.int32) # pylint: disable=protected-access + ) # 6. Configure Input Mapping def custom_gen_model_input_fn(batch): @@ -782,10 +819,21 @@ def custom_gen_model_input_fn(batch): ) return inputs_dict - trainer = trainer.with_gen_model_input_fn(custom_gen_model_input_fn) + if not diffusion_opd_enabled: + trainer = trainer.with_gen_model_input_fn(custom_gen_model_input_fn) # 7. Create Iterator Wrappers (Use Utils) train_iter = distillation_utils.MaxTextToTunixIterator(raw_train_iter) + if diffusion_opd_enabled: + assert diffusion_opd_lib is not None + train_iter = diffusion_opd_lib.replay_and_bound_iterator( + train_iter, + iter_steps=trainer.iter_steps, + train_steps=trainer.train_steps, + max_steps=student_config.steps, + accumulation_steps=student_config.gradient_accumulation_steps, + replay_iterator=getattr(raw_train_iter, "local_iterator", None), + ) eval_iter = None if raw_eval_iter is not None: @@ -821,6 +869,7 @@ def custom_gen_model_input_fn(batch): optimizer=trainer.optimizer, save_only_lora_params=getattr(trainer, "_lora_enabled", False), force=True, + custom_metadata=trainer.custom_checkpoint_metadata() if diffusion_opd_enabled else {}, ) if saved: # Ensure underlying orbax manager finishes writing @@ -911,6 +960,12 @@ def main(argv: Sequence[str]) -> None: teacher_argv = [argv[0], argv[1]] teacher_config = pyconfig.initialize(teacher_argv, **teacher_overrides) + if student_config.distill_data_source == "student_rollout": + # pylint: disable-next=import-outside-toplevel + from maxtext.trainers.post_train.distillation import diffusion_opd as diffusion_opd_lib + + diffusion_opd_lib.validate_diffusion_opd_configs(student_config, teacher_config, is_offline=is_offline) + # Batch shape (per_device_batch_size / max_target_length / gradient_accumulation_steps) # must be set at the YAML top level — not inside *_overrides — since student and # teacher share the input pipeline. diff --git a/tests/post_training/unit/maxtext_diffusion_opd_test.py b/tests/post_training/unit/maxtext_diffusion_opd_test.py new file mode 100644 index 0000000000..96ae731155 --- /dev/null +++ b/tests/post_training/unit/maxtext_diffusion_opd_test.py @@ -0,0 +1,539 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests MaxText preparation of fresh diffusion OPD batches.""" + +from types import SimpleNamespace +from unittest import mock + +from absl.testing import absltest +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.trainers.post_train.distillation import diffusion_opd +from maxtext.trainers.post_train.distillation import distillation_utils + + +def _student_config(**overrides): + """Builds the student configuration subset used by OPD validation.""" + values = { + "distill_data_source": "student_rollout", + "training_objective": "block_diffusion", + "attention_type": "block_diffusion", + "use_sft": True, + "sft_train_on_completion_only": True, + "vocab_size": 16, + "tokenizer_path": "test-tokenizer", + "tokenizer_type": "huggingface", + "distill_beta": 0.0, + "distill_beta_end": None, + "distill_alpha": 0.75, + "distill_alpha_end": None, + "distill_alpha_schedule": "constant", + "distill_temperature": 2.0, + "distill_temperature_end": None, + "distill_temperature_schedule": "constant", + "eval_interval": -1, + "steps": 10, + "gradient_accumulation_steps": 1, + "use_tunix_gradient_accumulation": False, + "max_target_length": 16, + "global_batch_size_to_load": 8, + "global_batch_size_to_train_on": 8, + "dataset_type": "hf", + "generate_padding_batch_train": False, + "elastic_enabled": False, + "expansion_factor_real_data": -1.0, + "enable_rampup_batch_size": False, + "learn_to_init_mode": False, + "student_params_to_update": None, + "distill_rollout_max_denoise_steps": -1, + "block_diffusion_block_size": 4, + "block_diffusion_mask_id": 15, + "block_diffusion_min_noise": 0.001, + "block_diffusion_logit_alignment": "same_position", + "block_diffusion_canvas_policy": "all_masked", + "distill_rollout_confidence_threshold": 0.9, + "distill_rollout_temperature": 1.0, + "distill_rollout_algorithm": "low_confidence", + "distill_rollout_stop_token_ids": [], + "enable_dropout": True, + "dropout_rate": 0.0, + "opt_type": "adamw", + "learning_rate": 1e-4, + "warmup_steps_fraction": 0.1, + "learning_rate_final_fraction": 0.1, + "gradient_clipping_threshold": 1.0, + "hardware": "cpu", + "shard_mode": "auto", + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _teacher_config(**overrides): + """Builds a causal teacher configuration subset.""" + values = { + "training_objective": "causal_lm", + "attention_type": "global", + "vocab_size": 16, + "tokenizer_path": "test-tokenizer", + "tokenizer_type": "huggingface", + "hardware": "cpu", + "shard_mode": "auto", + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _training_input(**overrides): + """Builds one explicit-mask target-aligned distillation batch.""" + values = { + "input_tokens": jnp.asarray([[3, 4, 31, 31]], dtype=jnp.int32), + "input_mask": jnp.ones((1, 4), dtype=jnp.bool_), + "positions": jnp.arange(4, dtype=jnp.int32)[None, :], + "decoder_segment_ids": jnp.ones((1, 4), dtype=jnp.int32), + "targets": jnp.asarray([[3, 4, 5, 6]], dtype=jnp.int32), + "targets_position": jnp.arange(4, dtype=jnp.int32)[None, :], + "targets_segmentation": jnp.ones((1, 4), dtype=jnp.int32), + "completion_mask": jnp.asarray([[0, 0, 1, 1]], dtype=jnp.int32), + "corruption_mask": jnp.asarray([[0, 0, 1, 1]], dtype=jnp.int32), + "targets_loss_mask": jnp.asarray([[0, 0, 1, 1]], dtype=jnp.int32), + } + values.update(overrides) + return distillation_utils.MaxTextTrainingInput(**values) + + +class MaxTextDiffusionOPDTest(absltest.TestCase): + + def test_config_validation_is_default_off(self): + student = _student_config(distill_data_source="dataset") + + self.assertFalse(diffusion_opd.validate_diffusion_opd_configs(student, _teacher_config(), is_offline=False)) + + def test_config_validation_accepts_supported_contract(self): + self.assertTrue(diffusion_opd.validate_diffusion_opd_configs(_student_config(), _teacher_config(), is_offline=False)) + + def test_config_validation_rejects_unsafe_modes(self): + cases = [ + (_student_config(), _teacher_config(), True, "online teacher"), + (_student_config(training_objective="causal_lm"), _teacher_config(), False, "student training_objective"), + (_student_config(use_sft=False), _teacher_config(), False, "completion-only SFT"), + (_student_config(), _teacher_config(training_objective="block_diffusion"), False, "causal teacher"), + (_student_config(), _teacher_config(attention_type="full"), False, "causal teacher"), + (_student_config(), _teacher_config(mtp_num_layers=1), False, "MTP teacher"), + (_student_config(), _teacher_config(tokenizer_path="other-tokenizer"), False, "identical.*tokenizer"), + ( + _student_config(tokenizer_type="sentencepiece"), + _teacher_config(tokenizer_type="sentencepiece"), + False, + "tokenizer_type='huggingface'", + ), + (_student_config(distill_beta=0.1), _teacher_config(), False, "feature distillation"), + (_student_config(distill_alpha=float("nan")), _teacher_config(), False, "finite"), + (_student_config(distill_rollout_temperature=float("inf")), _teacher_config(), False, "finite"), + (_student_config(eval_interval=10), _teacher_config(), False, "evaluation"), + ( + _student_config(generate_padding_batch_train=True), + _teacher_config(), + False, + "generate_padding_batch_train", + ), + (_student_config(elastic_enabled=True), _teacher_config(), False, "elastic data loading"), + (_student_config(expansion_factor_real_data=2.0), _teacher_config(), False, "expansion_factor_real_data"), + (_student_config(enable_rampup_batch_size=True), _teacher_config(), False, "ramp-up batch sizing"), + ( + _student_config(gradient_accumulation_steps=2), + _teacher_config(), + False, + "use_tunix_gradient_accumulation=True", + ), + (_student_config(distill_rollout_max_denoise_steps=3), _teacher_config(), False, "at least"), + ( + _student_config(distill_rollout_stop_token_ids=[7, 7]), + _teacher_config(), + False, + "unique IDs", + ), + ] + for student, teacher, is_offline, message in cases: + with self.subTest(message=message), self.assertRaisesRegex(ValueError, message): + diffusion_opd.validate_diffusion_opd_configs(student, teacher, is_offline=is_offline) + + def test_tpu_auto_sharding_is_preserved_for_models_without_explicit_support(self): + self.assertTrue( + diffusion_opd.validate_diffusion_opd_configs( + _student_config(hardware="tpu", shard_mode="auto"), + _teacher_config(hardware="tpu", shard_mode="auto"), + is_offline=False, + ) + ) + + def test_causal_teacher_logits_are_shifted_to_physical_targets(self): + raw_logits = jnp.arange(8, dtype=jnp.float32).reshape(1, 4, 2) + model = mock.Mock(return_value=raw_logits) + logits_fn = diffusion_opd.create_target_aligned_logits_fn(_teacher_config(), enable_dropout=False) + model_inputs = { + "input_tokens": jnp.ones((1, 4), dtype=jnp.int32), + "positions": jnp.arange(4, dtype=jnp.int32)[None, :], + "decoder_segment_ids": jnp.ones((1, 4), dtype=jnp.int32), + "targets": jnp.ones((1, 4), dtype=jnp.int32), + "targets_segmentation": jnp.ones((1, 4), dtype=jnp.int32), + } + + aligned = logits_fn(model, model_inputs) + + np.testing.assert_array_equal(aligned, raw_logits[:, [0, 0, 1, 2], :]) + self.assertFalse(model.call_args.kwargs["enable_dropout"]) + + def test_compiled_rollout_functionalizes_mutating_nnx_model(self): + class MutatingLogitModel(nnx.Module): + + def __init__(self): + self.counter = nnx.Variable(jnp.asarray(0, dtype=jnp.int32)) + + def __call__(self, decoder_input_tokens, **kwargs): + del kwargs + self.counter[...] += 1 + targets = jnp.full_like(decoder_input_tokens, 7) + return jax.nn.one_hot(targets, 16, dtype=jnp.float32) * 12.0 + + model = MutatingLogitModel() + rollout_fn = diffusion_opd.create_rollout_fn(_student_config()) + + generated = rollout_fn( + model, + jnp.asarray([[3, 4, 1, 1]], dtype=jnp.int32), + jnp.arange(4, dtype=jnp.int32)[None, :], + jnp.ones((1, 4), dtype=jnp.int32), + jnp.asarray([[0, 0, 1, 1]], dtype=jnp.bool_), + ) + + np.testing.assert_array_equal(generated, [[3, 4, 7, 7]]) + self.assertEqual(int(model.counter[...]), 0) + + def test_preparation_scores_clean_rollout_and_corrupts_only_student(self): + generated = jnp.asarray([[3, 4, 7, 8]], dtype=jnp.int32) + rollout_fn = mock.Mock(return_value=generated) + teacher_model = mock.sentinel.teacher_model + + def teacher_score_fn(model, model_inputs): + self.assertIs(model, teacher_model) + np.testing.assert_array_equal(model_inputs["input_tokens"], generated) + return jax.nn.one_hot(generated, 16, dtype=jnp.float32) + + batch = diffusion_opd.prepare_diffusion_opd_batch( + _training_input(), + mock.sentinel.student_model, + teacher_model, + rollout_fn=rollout_fn, + teacher_score_fn=teacher_score_fn, + mask_id=15, + ) + + np.testing.assert_array_equal(batch.student_batch.target_ids, generated) + np.testing.assert_array_equal(batch.student_batch.model_inputs["input_tokens"], [[3, 4, 15, 15]]) + np.testing.assert_array_equal(batch.student_batch.loss_weights, [[0.0, 0.0, 1.0, 1.0]]) + self.assertEqual(batch.teacher_logits.shape, (1, 4, 16)) + rollout_fn.assert_called_once() + + def test_preparation_excludes_tokens_after_generated_eos(self): + generated = jnp.asarray([[3, 4, 7, 8]], dtype=jnp.int32) + teacher_score_fn = mock.Mock(return_value=jax.nn.one_hot(generated, 16)) + + batch = diffusion_opd.prepare_diffusion_opd_batch( + _training_input(), + mock.sentinel.student_model, + mock.sentinel.teacher_model, + rollout_fn=mock.Mock(return_value=generated), + teacher_score_fn=teacher_score_fn, + mask_id=15, + stop_token_ids=(7,), + ) + + np.testing.assert_array_equal(batch.student_batch.loss_weights, [[0.0, 0.0, 1.0, 0.0]]) + np.testing.assert_array_equal(batch.student_batch.model_inputs["decoder_segment_ids"], [[1, 1, 1, 0]]) + np.testing.assert_array_equal(teacher_score_fn.call_args.args[1]["targets_segmentation"], [[1, 1, 1, 0]]) + + def test_shifted_block_anchor_can_be_weighted_without_corruption(self): + generated = jnp.asarray([[3, 7, 8, 9]], dtype=jnp.int32) + batch = diffusion_opd.prepare_diffusion_opd_batch( + _training_input( + completion_mask=jnp.asarray([[0, 1, 1, 1]], dtype=jnp.int32), + corruption_mask=jnp.asarray([[0, 1, 0, 1]], dtype=jnp.int32), + targets_loss_mask=jnp.asarray([[0, 1, 1, 1]], dtype=jnp.int32), + ), + mock.sentinel.student, + mock.sentinel.teacher, + rollout_fn=mock.Mock(return_value=generated), + teacher_score_fn=mock.Mock(return_value=jax.nn.one_hot(generated, 16)), + mask_id=15, + logit_alignment="shifted", + block_size=2, + ) + + np.testing.assert_array_equal(batch.student_batch.model_inputs["input_tokens"], [[3, 15, 8, 15]]) + np.testing.assert_array_equal(batch.student_batch.loss_weights, [[0.0, 1.0, 1.0, 1.0]]) + + def test_prepared_batch_runs_weighted_tunix_loss_and_student_gradient(self): + class TrainableLogitModel(nnx.Module): + + def __init__(self): + self.bias = nnx.Param(jnp.zeros((16,), dtype=jnp.float32)) + + def __call__(self, decoder_input_tokens, **kwargs): + del kwargs + return jnp.broadcast_to(self.bias[...], (*decoder_input_tokens.shape, 16)) + + generated = jnp.asarray([[3, 4, 7, 8]], dtype=jnp.int32) + batch = diffusion_opd.prepare_diffusion_opd_batch( + _training_input(), + mock.sentinel.student, + mock.sentinel.teacher, + rollout_fn=mock.Mock(return_value=generated), + teacher_score_fn=mock.Mock(return_value=jax.nn.one_hot(generated, 16) * 4.0), + mask_id=15, + ) + model = TrainableLogitModel() + logits_fn = diffusion_opd.create_target_aligned_logits_fn(_student_config(), enable_dropout=False) + + def loss_fn(trainable_model): + output = diffusion_opd.tunix_diffusion_opd.diffusion_opd_loss_fn( + trainable_model, + batch, + logits_fn, + temperature=2.0, + soft_loss_weight=0.75, + hard_loss_weight=0.25, + ) + return output.primary_loss.unreduced_sum, output + + (loss_sum, output), gradients = nnx.value_and_grad(loss_fn, has_aux=True)(model) + + self.assertTrue(bool(jnp.isfinite(loss_sum))) + self.assertEqual(float(output.primary_loss.denominator), 2.0) + self.assertTrue(any(bool(jnp.any(leaf[...] != 0)) for leaf in jax.tree.leaves(gradients))) + + def test_same_position_rejects_weighted_clean_targets(self): + with self.assertRaisesRegex(ValueError, "same-position targets must be corrupted"): + diffusion_opd.prepare_diffusion_opd_batch( + _training_input( + corruption_mask=jnp.asarray([[0, 0, 1, 0]], dtype=jnp.int32), + targets_loss_mask=jnp.asarray([[0, 0, 1, 1]], dtype=jnp.int32), + ), + mock.sentinel.student, + mock.sentinel.teacher, + rollout_fn=mock.Mock(), + teacher_score_fn=mock.Mock(), + mask_id=15, + logit_alignment="same_position", + block_size=4, + ) + + def test_device_safe_weights_zero_invalid_suffix_rows(self): + weights = diffusion_opd._safe_loss_weights( # pylint: disable=protected-access + jnp.arange(4, dtype=jnp.int32)[None, :], + jnp.ones((1, 4), dtype=jnp.bool_), + jnp.asarray([[0, 1, 0, 1]], dtype=jnp.bool_), + jnp.asarray([[0, 1, 0, 1]], dtype=jnp.bool_), + jnp.asarray([[0.0, 1.0, 0.0, 1.0]], dtype=jnp.float32), + alignment="same_position", + block_size=4, + ) + + np.testing.assert_array_equal(weights, jnp.zeros((1, 4), dtype=jnp.float32)) + + def test_device_safe_weights_zero_out_of_range_positions(self): + for positions in (jnp.asarray([[-1, 1]]), jnp.asarray([[0, 2]])): + with self.subTest(positions=positions): + weights = diffusion_opd._safe_loss_weights( # pylint: disable=protected-access + positions, + jnp.ones((1, 2), dtype=jnp.bool_), + jnp.asarray([[0, 1]], dtype=jnp.bool_), + jnp.asarray([[0, 1]], dtype=jnp.bool_), + jnp.asarray([[0.0, 1.0]], dtype=jnp.float32), + alignment="same_position", + block_size=2, + ) + + np.testing.assert_array_equal(weights, jnp.zeros((1, 2), dtype=jnp.float32)) + + def test_device_safe_weights_zero_mask_subset_violations(self): + weights = diffusion_opd._safe_loss_weights( # pylint: disable=protected-access + jnp.asarray([[0, 1, 2]]), + jnp.asarray([[1, 1, 0]], dtype=jnp.bool_), + jnp.asarray([[0, 1, 1]], dtype=jnp.bool_), + jnp.asarray([[0, 1, 1]], dtype=jnp.bool_), + jnp.asarray([[0.0, 1.0, 1.0]], dtype=jnp.float32), + alignment="same_position", + block_size=2, + ) + + np.testing.assert_array_equal(weights, jnp.zeros((1, 3), dtype=jnp.float32)) + + def test_preparation_rejects_missing_or_cross_role_masks(self): + rollout_fn = mock.Mock() + teacher_score_fn = mock.Mock() + with self.assertRaisesRegex(ValueError, "corruption_mask"): + diffusion_opd.prepare_diffusion_opd_batch( + _training_input(corruption_mask=None), + mock.sentinel.student, + mock.sentinel.teacher, + rollout_fn=rollout_fn, + teacher_score_fn=teacher_score_fn, + mask_id=15, + ) + with self.assertRaisesRegex(ValueError, "subset of completion_mask"): + diffusion_opd.prepare_diffusion_opd_batch( + _training_input(corruption_mask=jnp.asarray([[0, 1, 1, 0]], dtype=jnp.int32)), + mock.sentinel.student, + mock.sentinel.teacher, + rollout_fn=rollout_fn, + teacher_score_fn=teacher_score_fn, + mask_id=15, + ) + rollout_fn.assert_not_called() + teacher_score_fn.assert_not_called() + + def test_hf_resume_replays_and_bounds_deterministic_stream(self): + bounded = diffusion_opd.replay_and_bound_iterator( + iter(range(20)), + iter_steps=4, + train_steps=2, + max_steps=5, + accumulation_steps=2, + ) + + self.assertEqual(list(bounded), [4, 5, 6, 7, 8, 9]) + + def test_hf_resume_can_replay_before_global_materialization(self): + local_iterator = iter(range(20)) + + class GlobalIterator: + + def __next__(self): + return next(local_iterator) * 10 + + def __iter__(self): + return self + + bounded = diffusion_opd.replay_and_bound_iterator( + GlobalIterator(), + iter_steps=4, + train_steps=2, + max_steps=5, + accumulation_steps=2, + replay_iterator=local_iterator, + ) + + self.assertEqual(list(bounded), [40, 50, 60, 70, 80, 90]) + + def test_hf_resume_rejects_early_exhaustion_during_replay(self): + with self.assertRaisesRegex(ValueError, "replay exhausted after 2 of 4"): + diffusion_opd.replay_and_bound_iterator( + iter(range(2)), + iter_steps=4, + train_steps=2, + max_steps=5, + accumulation_steps=2, + ) + + def test_hf_resume_rejects_early_exhaustion_during_remaining_training(self): + bounded = diffusion_opd.replay_and_bound_iterator( + iter(range(4)), + iter_steps=4, + train_steps=2, + max_steps=5, + accumulation_steps=2, + replay_iterator=iter(range(4)), + ) + + with self.assertRaisesRegex(ValueError, "4 of 6 required remaining microbatches"): + list(bounded) + + def test_hf_resume_rejects_checkpoint_beyond_configured_steps(self): + with self.assertRaisesRegex(ValueError, "exceeds configured max_steps"): + diffusion_opd.replay_and_bound_iterator( + iter(()), + iter_steps=12, + train_steps=6, + max_steps=5, + accumulation_steps=2, + ) + + def test_hf_resume_at_final_step_does_not_replay(self): + bounded = diffusion_opd.replay_and_bound_iterator( + iter(()), + iter_steps=10, + train_steps=5, + max_steps=5, + accumulation_steps=2, + ) + + self.assertEqual(list(bounded), []) + + def test_checkpoint_contract_captures_batch_and_preprocessing_identity(self): + base = diffusion_opd._checkpoint_contract(_student_config(), _teacher_config()) # pylint: disable=protected-access + changed = diffusion_opd._checkpoint_contract( # pylint: disable=protected-access + _student_config(gradient_accumulation_steps=2, use_tunix_gradient_accumulation=True), _teacher_config() + ) + + self.assertNotEqual(base, changed) + self.assertEqual(base["max_target_length"], 16) + + def test_checkpoint_restore_preserves_contract_without_claiming_hf_cursor_state(self): + trainer = diffusion_opd.MaxTextDiffusionOPDTrainer.__new__(diffusion_opd.MaxTextDiffusionOPDTrainer) + trainer.checkpoint_manager = mock.Mock() + trainer.model = mock.sentinel.model + trainer.optimizer = mock.sentinel.optimizer + trainer.config = mock.Mock(checkpointing_options=mock.sentinel.options) + trainer.config.get_with_default.return_value = 1 + trainer._lora_enabled = False # pylint: disable=protected-access + trainer.checkpoint_contract = {"version": 1, "teacher_model": "teacher"} + manager = mock.Mock() + manager.maybe_restore.return_value = (4, {"diffusion_opd_contract": trainer.checkpoint_contract}) + raw_iterator = SimpleNamespace(local_iterator=mock.sentinel.local_iterator) + + with mock.patch.object(diffusion_opd.distillation_utils, "MaxTextCheckpointManager", return_value=manager) as factory: + result = trainer.setup_checkpoint_manager_and_restore( + raw_iterator, SimpleNamespace(dataset_type="hf", checkpoint_dir="/tmp") + ) + + self.assertIs(result, raw_iterator) + self.assertIsNone(factory.call_args.kwargs["raw_iterator"]) + self.assertEqual(trainer._train_steps, 4) # pylint: disable=protected-access + self.assertEqual(trainer.custom_checkpoint_metadata()["diffusion_opd_contract"], trainer.checkpoint_contract) + + def test_checkpoint_restore_rejects_semantic_contract_mismatch(self): + trainer = diffusion_opd.MaxTextDiffusionOPDTrainer.__new__(diffusion_opd.MaxTextDiffusionOPDTrainer) + trainer.checkpoint_manager = mock.Mock() + trainer.model = mock.sentinel.model + trainer.optimizer = mock.sentinel.optimizer + trainer.config = mock.Mock(checkpointing_options=mock.sentinel.options) + trainer.config.get_with_default.return_value = 1 + trainer._lora_enabled = False # pylint: disable=protected-access + trainer.checkpoint_contract = {"version": 1, "teacher_model": "expected"} + manager = mock.Mock() + manager.maybe_restore.return_value = (4, {"diffusion_opd_contract": {"version": 1, "teacher_model": "other"}}) + + with ( + mock.patch.object(diffusion_opd.distillation_utils, "MaxTextCheckpointManager", return_value=manager), + self.assertRaisesRegex(ValueError, "checkpoint contract"), + ): + trainer.setup_checkpoint_manager_and_restore( + SimpleNamespace(local_iterator=mock.sentinel.local_iterator), + SimpleNamespace(dataset_type="hf", checkpoint_dir="/tmp"), + ) diff --git a/tests/post_training/unit/train_distill_test.py b/tests/post_training/unit/train_distill_test.py index 24cd1d0bc8..6a13c80ef0 100644 --- a/tests/post_training/unit/train_distill_test.py +++ b/tests/post_training/unit/train_distill_test.py @@ -104,6 +104,9 @@ def test_maxtext_to_tunix_iterator_sft(self): "targets": np.array([[11, 12]]), "targets_position": np.array([[100, 101]]), # Custom position "targets_segmentation": np.array([[0, 1]]), # Custom segmentation (mask) + "completion_mask": np.array([[0, 1]]), + "corruption_mask": np.array([[0, 1]]), + "targets_loss_mask": np.array([[0, 1]]), } dummy_iter_sft = iter([dummy_batch_sft]) @@ -115,6 +118,9 @@ def test_maxtext_to_tunix_iterator_sft(self): self.assertIsInstance(tunix_input_sft, distillation_utils.MaxTextTrainingInput) np.testing.assert_array_equal(tunix_input_sft.targets_position, dummy_batch_sft["targets_position"]) np.testing.assert_array_equal(tunix_input_sft.targets_segmentation, dummy_batch_sft["targets_segmentation"]) + np.testing.assert_array_equal(tunix_input_sft.completion_mask, dummy_batch_sft["completion_mask"]) + np.testing.assert_array_equal(tunix_input_sft.corruption_mask, dummy_batch_sft["corruption_mask"]) + np.testing.assert_array_equal(tunix_input_sft.targets_loss_mask, dummy_batch_sft["targets_loss_mask"]) def test_maxtext_to_tunix_iterator_packed_fallback(self): """Verifies fallback behavior when segmentation is missing.""" @@ -151,13 +157,19 @@ def test_prepare_inputs_logic(self): input_mask=jnp.array([[True]]), positions=jnp.array([[0]]), targets=jnp.array([[1]]), + completion_mask=jnp.array([[1]]), + corruption_mask=jnp.array([[1]]), + targets_loss_mask=jnp.array([[1]]), ) # 4. Run - _ = trainer._prepare_inputs(input_data) + prepared = trainer._prepare_inputs(input_data) # 5. Verify trainer.strategy.get_teacher_outputs.assert_not_called() + np.testing.assert_array_equal(prepared.completion_mask, input_data.completion_mask) + np.testing.assert_array_equal(prepared.corruption_mask, input_data.corruption_mask) + np.testing.assert_array_equal(prepared.targets_loss_mask, input_data.targets_loss_mask) @mock.patch("maxtext.trainers.post_train.distillation.train_distill.optax.global_norm") @mock.patch("maxtext.trainers.post_train.distillation.train_distill.jax.tree.map") @@ -646,6 +658,32 @@ def test_setup_pipeline_restored(self): # Verify it returned the restored iterator, NOT the raw one self.assertEqual(result, restored_iter) + def test_standard_distillation_rejects_diffusion_opd_checkpoint(self): + mock_trainer = mock.Mock() + mock_trainer.checkpoint_manager = mock.Mock() + mock_trainer.model = mock.Mock() + mock_trainer.optimizer = mock.Mock() + mock_trainer._lora_enabled = False + mock_trainer.config = mock.Mock( + checkpointing_options=ocp.CheckpointManagerOptions(max_to_keep=1, create=True) + ) + mock_manager = mock.Mock() + mock_manager.maybe_restore.return_value = (10, {"diffusion_opd_contract": {"version": 1}}) + config = mock.Mock(dataset_type="synthetic", checkpoint_dir=self.test_dir) + + with ( + mock.patch( + "maxtext.trainers.post_train.distillation.distillation_utils.MaxTextCheckpointManager", + return_value=mock_manager, + ), + self.assertRaisesRegex(ValueError, "cannot resume a diffusion OPD checkpoint"), + ): + train_distill.MaxTextDistillationTrainer.setup_checkpoint_manager_and_restore( + mock_trainer, + object(), + config, + ) + def test_eval_step_calls_student_forward(self): """Verifies eval step correctly calls the student forward function and computes loss.""" # 1. Initialize Trainer (bypass init) diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 32a4bbba0e..a4e8abad40 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -241,6 +241,9 @@ def test_default_attention_remains_global(self): self.assertEqual(config.attention_type, "global") self.assertEqual(config.training_objective, "causal_lm") self.assertEqual(config.block_diffusion_mask_id, -1) + self.assertEqual(config.distill_data_source, "dataset") + self.assertEqual(config.distill_rollout_algorithm, "low_confidence") + self.assertEqual(config.distill_rollout_stop_token_ids, []) def test_block_diffusion_training_objective_config(self): argv = [ diff --git a/tests/unit/input_pipeline_utils_test.py b/tests/unit/input_pipeline_utils_test.py index fe5ce30cda..ec775a949e 100644 --- a/tests/unit/input_pipeline_utils_test.py +++ b/tests/unit/input_pipeline_utils_test.py @@ -241,6 +241,77 @@ def test_completion_only_requires_explicit_role_mask(self): with self.assertRaisesRegex(ValueError, "explicit completion_mask"): self._apply_transform(clean) + def test_opd_suffix_contract_selects_final_assistant_turn(self): + clean = _make_block_diffusion_batch() + clean["completion_mask"][0] = np.asarray([0, 0, 1, 1, 0, 0, 1, 1, 1, 0]) + transform = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + select_last_completion_suffix=True, + ) + + output = transform.random_map(clean, np.random.default_rng(0)) + + np.testing.assert_array_equal( + output["completion_mask"][0], + np.asarray([0, 0, 0, 0, 0, 0, 1, 1, 1, 0]), + ) + + def test_opd_suffix_contract_truncates_context_after_last_completion(self): + clean = _make_block_diffusion_batch() + clean["completion_mask"][0] = np.asarray([0, 0, 1, 1, 0, 0, 0, 0, 0, 0]) + transform = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + select_last_completion_suffix=True, + ) + + output = transform.random_map(clean, np.random.default_rng(0)) + + np.testing.assert_array_equal( + output["completion_mask"][0], + np.asarray([0, 0, 1, 1, 0, 0, 0, 0, 0, 0]), + ) + np.testing.assert_array_equal( + output["inputs_segmentation"][0], + np.asarray([1, 1, 1, 1, 0, 0, 0, 0, 0, 0]), + ) + np.testing.assert_array_equal( + output["targets_segmentation"][0], + np.asarray([1, 1, 1, 1, 0, 0, 0, 0, 0, 0]), + ) + np.testing.assert_array_equal(output["targets"][0, 4:], 0) + + def test_opd_suffix_contract_requires_prompt_and_completion(self): + transform = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + select_last_completion_suffix=True, + ) + for completion_mask in ( + np.zeros(10, dtype=np.int32), + np.asarray([1, 1, 1, 1, 1, 1, 1, 1, 1, 0]), + ): + clean = _make_block_diffusion_batch() + clean["completion_mask"][0] = completion_mask + with self.assertRaisesRegex(ValueError, "completion span after a prompt"): + transform.random_map(clean, np.random.default_rng(0)) + + def test_opd_suffix_contract_accepts_prompt_then_completion(self): + transform = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + select_last_completion_suffix=True, + ) + + output = transform.random_map(_make_block_diffusion_batch(), np.random.default_rng(0)) + + self.assertTrue(output["targets_loss_mask"].any()) + def test_completion_mask_does_not_replace_full_validity(self): clean = _make_block_diffusion_batch() transform = BlockDiffusionCorruption(block_size=4, mask_id=99, min_noise=1.0, completion_only=False) From 6ec8cce3ede20d9a72ec00d3045b37d204482cc9 Mon Sep 17 00:00:00 2001 From: ethannnnnn Date: Thu, 23 Jul 2026 02:19:39 -0700 Subject: [PATCH 5/6] [maxtext] Export exact stochastic denoising traces Extend the low-confidence block denoiser with a stochastic, per-row rollout that records the sampled token, action step, and action log probability for every completion position. The trace keeps the full fixed-horizon denoising schedule even after a visible EOS, isolates RNG streams between batch rows, excludes the mask token from sampling, and forces progress when no token clears the confidence threshold. Existing deterministic generation remains unchanged. This trace is the minimal model-independent contract needed to replay the same diffusion actions under live, reference, and old policies during policy optimization. Test Plan: - pytest tests/unit/diffusion_denoise_test.py - included in the 83-test focused diffusion-RL suite - Pyink, Pylint, py_compile, and git diff --check --- src/maxtext/diffusion/denoise.py | 184 ++++++++++++++++++++++---- tests/unit/diffusion_denoise_test.py | 189 +++++++++++++++++++++++++++ 2 files changed, 351 insertions(+), 22 deletions(-) diff --git a/src/maxtext/diffusion/denoise.py b/src/maxtext/diffusion/denoise.py index 2f99b91069..937c146a49 100644 --- a/src/maxtext/diffusion/denoise.py +++ b/src/maxtext/diffusion/denoise.py @@ -15,12 +15,21 @@ """Model-independent block-diffusion rollout state transitions.""" from collections.abc import Callable +from typing import NamedTuple import jax import jax.numpy as jnp import numpy as np +class DenoiseTrace(NamedTuple): + """Compact sampled trajectory for replaying diffusion policy scores.""" + + tokens: jax.Array + action_steps: jax.Array + action_logps: jax.Array + + def _validate_shapes(initial_tokens, positions, validity_mask, completion_mask): """Checks that all token-level rollout arrays share a batch-major shape.""" expected_shape = tuple(initial_tokens.shape) @@ -43,6 +52,26 @@ def _concrete_numpy(value): return np.asarray(value) +def _per_row_rngs(rng, batch_size): + """Normalizes one or per-row typed/legacy JAX keys to per-row keys.""" + rng = jnp.asarray(rng) + is_typed_key = jax.dtypes.issubdtype(rng.dtype, jax.dtypes.prng_key) + scalar_shape = () if is_typed_key else tuple(jax.random.PRNGKey(0).shape) + if tuple(rng.shape) == scalar_shape: + return jax.vmap(lambda row: jax.random.fold_in(rng, row))(jnp.arange(batch_size, dtype=jnp.int32)) + if tuple(rng.shape) == (batch_size, *scalar_shape): + return rng + raise ValueError( + "rng must be one PRNG key or one key per batch row; " + f"received shape {tuple(rng.shape)} for key shape {scalar_shape}" + ) + + +def _select_row_rngs(row_mask, new_rngs, old_rngs): + selector = row_mask.reshape((row_mask.shape[0],) + (1,) * (old_rngs.ndim - 1)) + return jnp.where(selector, new_rngs, old_rngs) + + def _validate_logical_positions(positions, validity_mask, completion_mask, *, shifted_seed): """Checks logical sequence invariants on eager, host-addressable inputs.""" concrete_positions = _concrete_numpy(positions) @@ -94,7 +123,7 @@ def validate_completion_suffix(positions, validity_mask, completion_mask, *, shi raise ValueError("diffusion OPD rollout requires completion_mask to be a contiguous suffix") -def low_confidence_generate( +def low_confidence_rollout( logits_fn: Callable[[jax.Array], jax.Array], initial_tokens: jax.Array, positions: jax.Array, @@ -108,13 +137,16 @@ def low_confidence_generate( confidence_threshold: float = 0.9, temperature: float = 1.0, max_denoise_steps: int | None = None, -) -> jax.Array: - """Generates a completion block by block with confidence-based commits. + rng: jax.Array | None = None, +) -> DenoiseTrace: + """Samples a completion and records each token's pre-commit action step. ``logits_fn`` must return target-aligned logits for the current token canvas. Each denoising step commits every token at or above the confidence threshold; if a row has no such token, its highest-confidence unresolved token is forced - to commit. This guarantees progress for heterogeneous and partial blocks. + to commit. ``action_steps`` compactly identifies the exact masked canvas for + replay: token ``p`` was scored while every token with step >= its step was + still masked. Prompt, padding, and inactive completion tokens use step -1. """ _validate_shapes(initial_tokens, positions, validity_mask, completion_mask) valid_contracts = {("same_position", "all_masked"), ("shifted", "seed_and_mask")} @@ -143,10 +175,15 @@ def low_confidence_generate( completion_mask = jnp.asarray(completion_mask, dtype=jnp.bool_) & validity_mask positions = jnp.asarray(positions, dtype=jnp.int32) canvas = jnp.where(completion_mask, jnp.asarray(mask_id, initial_tokens.dtype), initial_tokens) + action_steps = jnp.full(initial_tokens.shape, -1, dtype=jnp.int32) + action_logps = jnp.zeros(initial_tokens.shape, dtype=jnp.float32) + sample_tokens = rng is not None + rng = jax.random.PRNGKey(0) if rng is None else jnp.asarray(rng) + row_rngs = _per_row_rngs(rng, initial_tokens.shape[0]) block_ids = positions // block_size num_blocks = (initial_tokens.shape[1] + block_size - 1) // block_size - def propose(current_canvas): + def propose(current_canvas, proposal_keys): logits = logits_fn(current_canvas) expected_prefix = tuple(current_canvas.shape) if len(logits.shape) != 3 or tuple(logits.shape[:2]) != expected_prefix: @@ -161,30 +198,63 @@ def propose(current_canvas): raise ValueError(f"mask_id must satisfy 0 <= mask_id < vocab_size ({vocab_size}); received {mask_id}") scaled_logits = jnp.asarray(logits, dtype=jnp.float32) / temperature scaled_logits = scaled_logits.at[..., mask_id].set(-jnp.inf) - probabilities = jax.nn.softmax(scaled_logits, axis=-1) - return jnp.argmax(scaled_logits, axis=-1).astype(initial_tokens.dtype), jnp.max(probabilities, axis=-1) + log_probabilities = jax.nn.log_softmax(scaled_logits, axis=-1) + if sample_tokens: + proposed_tokens = jax.vmap(lambda key, row_logits: jax.random.categorical(key, row_logits, axis=-1))( + proposal_keys, scaled_logits + ).astype(initial_tokens.dtype) + else: + proposed_tokens = jnp.argmax(scaled_logits, axis=-1).astype(initial_tokens.dtype) + proposed_logps = jnp.take_along_axis(log_probabilities, proposed_tokens[..., None], axis=-1)[..., 0] + confidence = jnp.max(jnp.exp(log_probabilities), axis=-1) + return proposed_tokens, confidence, proposed_logps - def generate_block(block_id, current_canvas): + def generate_block(block_id, rollout_state): + current_canvas, current_steps, current_logps, global_step, current_rng = rollout_state in_block = completion_mask & (block_ids == block_id) - def run_active_block(active_canvas): + def run_active_block(active_state): + active_canvas, active_steps, active_logps, active_global_step, active_rng = active_state anchors = in_block & (positions % block_size == 0) if shifted_seed else jnp.zeros_like(in_block) - def generate_anchors(anchor_canvas): - anchor_tokens, _ = propose(anchor_canvas) - return jnp.where(anchors, anchor_tokens, anchor_canvas) + def generate_anchors(anchor_state): + anchor_canvas, anchor_steps, anchor_logps, anchor_global_step, anchor_rng = anchor_state + split_keys = jax.vmap(jax.random.split)(anchor_rng) + next_rng = split_keys[:, 0] + proposal_keys = split_keys[:, 1] + anchor_tokens, _, proposed_logps = propose(anchor_canvas, proposal_keys) + anchor_canvas = jnp.where(anchors, anchor_tokens, anchor_canvas) + anchor_steps = jnp.where(anchors, anchor_global_step[:, None], anchor_steps) + anchor_logps = jnp.where(anchors, proposed_logps, anchor_logps) + rows_with_anchors = jnp.any(anchors, axis=1) + anchor_rng = _select_row_rngs(rows_with_anchors, next_rng, anchor_rng) + return ( + anchor_canvas, + anchor_steps, + anchor_logps, + anchor_global_step + rows_with_anchors.astype(jnp.int32), + anchor_rng, + ) if shifted_seed: - active_canvas = jax.lax.cond(jnp.any(anchors), generate_anchors, lambda value: value, active_canvas) + active_canvas, active_steps, active_logps, active_global_step, active_rng = jax.lax.cond( + jnp.any(anchors), + generate_anchors, + lambda value: value, + (active_canvas, active_steps, active_logps, active_global_step, active_rng), + ) unresolved = in_block & ~anchors def continue_denoising(state): - step, _, remaining = state + step, _, _, _, _, _, remaining = state return (step < max_denoise_steps) & jnp.any(remaining) def denoise_step(state): - step, step_canvas, remaining = state - proposed_tokens, confidence = propose(step_canvas) + step, step_canvas, step_ids, step_logps, step_global, step_rng, remaining = state + split_keys = jax.vmap(jax.random.split)(step_rng) + next_rng = split_keys[:, 0] + proposal_keys = split_keys[:, 1] + proposed_tokens, confidence, proposed_logps = propose(step_canvas, proposal_keys) commits = remaining & (confidence >= confidence_threshold) row_needs_fallback = jnp.any(remaining, axis=1) & ~jnp.any(commits, axis=1) fallback_confidence = jnp.max(jnp.where(remaining, confidence, -jnp.inf), axis=1) @@ -194,15 +264,85 @@ def denoise_step(state): fallback = jax.nn.one_hot(fallback_indices, remaining.shape[1], dtype=jnp.bool_) commits |= fallback & row_needs_fallback[:, None] step_canvas = jnp.where(commits, proposed_tokens, step_canvas) - return step + 1, step_canvas, remaining & ~commits + step_ids = jnp.where(commits, step_global[:, None], step_ids) + step_logps = jnp.where(commits, proposed_logps, step_logps) + rows_with_commits = jnp.any(commits, axis=1) + step_rng = _select_row_rngs(rows_with_commits, next_rng, step_rng) + return ( + step + 1, + step_canvas, + step_ids, + step_logps, + step_global + rows_with_commits.astype(jnp.int32), + step_rng, + remaining & ~commits, + ) - _, active_canvas, _ = jax.lax.while_loop( + _, active_canvas, active_steps, active_logps, active_global_step, active_rng, _ = jax.lax.while_loop( continue_denoising, denoise_step, - (jnp.asarray(0, dtype=jnp.int32), active_canvas, unresolved), + ( + jnp.asarray(0, dtype=jnp.int32), + active_canvas, + active_steps, + active_logps, + active_global_step, + active_rng, + unresolved, + ), ) - return active_canvas + return active_canvas, active_steps, active_logps, active_global_step, active_rng + + return jax.lax.cond( + jnp.any(in_block), + run_active_block, + lambda value: value, + (current_canvas, current_steps, current_logps, global_step, current_rng), + ) - return jax.lax.cond(jnp.any(in_block), run_active_block, lambda value: value, current_canvas) + canvas, action_steps, action_logps, _, _ = jax.lax.fori_loop( + 0, + num_blocks, + generate_block, + ( + canvas, + action_steps, + action_logps, + jnp.zeros((initial_tokens.shape[0],), dtype=jnp.int32), + row_rngs, + ), + ) + return DenoiseTrace(tokens=canvas, action_steps=action_steps, action_logps=action_logps) + + +def low_confidence_generate( + logits_fn: Callable[[jax.Array], jax.Array], + initial_tokens: jax.Array, + positions: jax.Array, + validity_mask: jax.Array, + completion_mask: jax.Array, + *, + block_size: int, + mask_id: int, + logit_alignment: str, + canvas_policy: str, + confidence_threshold: float = 0.9, + temperature: float = 1.0, + max_denoise_steps: int | None = None, +) -> jax.Array: + """Greedily generates tokens while preserving the original public API.""" - return jax.lax.fori_loop(0, num_blocks, generate_block, canvas) + return low_confidence_rollout( + logits_fn, + initial_tokens, + positions, + validity_mask, + completion_mask, + block_size=block_size, + mask_id=mask_id, + logit_alignment=logit_alignment, + canvas_policy=canvas_policy, + confidence_threshold=confidence_threshold, + temperature=temperature, + max_denoise_steps=max_denoise_steps, + ).tokens diff --git a/tests/unit/diffusion_denoise_test.py b/tests/unit/diffusion_denoise_test.py index 37bf84491d..d00df168e5 100644 --- a/tests/unit/diffusion_denoise_test.py +++ b/tests/unit/diffusion_denoise_test.py @@ -28,6 +28,195 @@ def _target_logits(targets, vocab_size=32, high=12.0): class DiffusionDenoiseTest(absltest.TestCase): + def test_trace_records_each_pre_commit_step_and_logprob(self): + initial = jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + + trace = denoise.low_confidence_rollout( + lambda canvas: jnp.zeros((*canvas.shape, 8), dtype=jnp.float32), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + ) + + np.testing.assert_array_equal(trace.action_steps, [[-1, 0, 1, 2]]) + np.testing.assert_allclose(trace.action_logps[:, 1:], -np.log(7.0), rtol=1e-6) + self.assertFalse(bool(jnp.any(trace.tokens == 7))) + + def test_trace_records_shifted_anchor_as_an_action(self): + initial = jnp.asarray([[2, 3, 4, 5, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(6, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 0, 0, 0, 1, 1]], dtype=jnp.bool_) + + trace = denoise.low_confidence_rollout( + lambda canvas: _target_logits(jnp.full_like(canvas, 6), vocab_size=8), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="shifted", + canvas_policy="seed_and_mask", + ) + + np.testing.assert_array_equal(trace.action_steps, [[-1, -1, -1, -1, 0, 1]]) + np.testing.assert_array_equal(trace.tokens, [[2, 3, 4, 5, 6, 6]]) + + def test_sampled_trace_is_seeded_and_excludes_mask_token(self): + initial = jnp.asarray([[4, 1, 1, 1, 1, 1, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(8, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 1, 1, 1, 1, 1, 1, 1]], dtype=jnp.bool_) + + def rollout(seed): + return denoise.low_confidence_rollout( + lambda canvas: jnp.zeros((*canvas.shape, 32), dtype=jnp.float32), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + rng=jax.random.PRNGKey(seed), + ) + + first = rollout(7) + repeated = rollout(7) + different = rollout(8) + + np.testing.assert_array_equal(first.tokens, repeated.tokens) + self.assertFalse(bool(jnp.array_equal(first.tokens, different.tokens))) + self.assertFalse(bool(jnp.any(first.tokens == 31))) + + def test_rollout_accepts_typed_and_legacy_keys_for_each_prng_implementation(self): + initial = jnp.asarray([[4, 7, 7, 7], [3, 7, 7, 7]], dtype=jnp.int32) + positions = jnp.broadcast_to(jnp.arange(4, dtype=jnp.int32), initial.shape) + completion = positions > 0 + + def logits_fn(canvas): + token_ids = (jnp.arange(canvas.shape[1])[None, :] + 1) % 7 + return jax.nn.one_hot(jnp.broadcast_to(token_ids, canvas.shape), 8, dtype=jnp.float32) * 4.0 + + for implementation in ("threefry2x32", "unsafe_rbg"): + with ( + self.subTest(implementation=implementation), + jax.default_prng_impl(implementation), + ): + legacy_key = jax.random.PRNGKey(7) + typed_key = jax.random.key(7) + keys = { + "legacy_scalar": legacy_key, + "legacy_batched": jax.vmap(lambda row: jax.random.fold_in(legacy_key, row))(jnp.arange(2)), + "typed_scalar": typed_key, + "typed_batched": jax.vmap(lambda row: jax.random.fold_in(typed_key, row))(jnp.arange(2)), + } + traces = {} + for key_kind, key in keys.items(): + with self.subTest(implementation=implementation, key_kind=key_kind): + traces[key_kind] = denoise.low_confidence_rollout( + logits_fn, + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + rng=key, + ) + self.assertFalse(bool(jnp.any(traces[key_kind].tokens == 7))) + + for trace in traces.values(): + np.testing.assert_array_equal(trace.tokens, traces["legacy_scalar"].tokens) + np.testing.assert_array_equal(trace.action_steps, traces["legacy_scalar"].action_steps) + + generated = denoise.low_confidence_generate( + logits_fn, + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + ) + self.assertFalse(bool(jnp.any(generated == 7))) + + def test_trace_steps_are_compact_per_row_with_heterogeneous_prompts(self): + initial = jnp.asarray( + [ + [0, 0, 0, 4, 1, 1, 1, 1], + [4, 5, 6, 7, 1, 1, 1, 1], + ], + dtype=jnp.int32, + ) + validity = initial != 0 + positions = jnp.where(validity, jnp.cumsum(validity, axis=1) - 1, 0) + completion = jnp.asarray( + [ + [0, 0, 0, 0, 1, 1, 1, 1], + [0, 0, 0, 0, 1, 1, 1, 1], + ], + dtype=jnp.bool_, + ) + + trace = denoise.low_confidence_rollout( + lambda canvas: jnp.zeros((*canvas.shape, 16), dtype=jnp.float32), + initial, + positions, + validity, + completion, + block_size=4, + mask_id=15, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + ) + + np.testing.assert_array_equal(trace.action_steps[:, 4:], [[0, 1, 2, 3], [0, 1, 2, 3]]) + + def test_row_sampling_is_independent_of_other_rows_denoise_progress(self): + first = jnp.asarray([[0, 0, 0, 4, 1, 1, 1, 1]], dtype=jnp.int32) + second = jnp.asarray([[4, 5, 6, 7, 1, 1, 1, 1]], dtype=jnp.int32) + + def rollout(initial): + validity = initial != 0 + positions = jnp.where(validity, jnp.cumsum(validity, axis=1) - 1, 0) + completion = jnp.zeros_like(validity).at[:, 4:].set(True) + return denoise.low_confidence_rollout( + lambda canvas: jnp.zeros((*canvas.shape, 16), dtype=jnp.float32), + initial, + positions, + validity, + completion, + block_size=4, + mask_id=15, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + rng=jax.random.PRNGKey(11), + ) + + alone = rollout(first) + batched = rollout(jnp.concatenate([first, second], axis=0)) + + np.testing.assert_array_equal(alone.tokens[0], batched.tokens[0]) + np.testing.assert_allclose(alone.action_logps[0], batched.action_logps[0]) + def test_same_position_generates_partial_blocks_and_preserves_prompt(self): initial = jnp.asarray([[7, 8, 1, 1, 1, 1, 1, 0]], dtype=jnp.int32) expected = jnp.asarray([[7, 8, 12, 13, 14, 15, 16, 0]], dtype=jnp.int32) From 90b5fc8d64f60068461ab014ee59887ece7b26fb Mon Sep 17 00:00:00 2001 From: ethannnnnn Date: Thu, 23 Jul 2026 02:19:46 -0700 Subject: [PATCH 6/6] [maxtext] Add exact block-diffusion policy replay Add a correctness-first MaxText/Tunix adapter for block-diffusion policy optimization. The in-process rollout exports the full stochastic denoising trace while keeping user-visible completions truncated at the first stop token. The scorer reconstructs each pre-commit canvas and evaluates the recorded action at its original denoising step, so rollout, live, reference, and old-policy log probabilities share one target-aligned contract. Generation uses independent train/eval RNG streams with resumable step contexts. Invalid stop IDs, non-finite logits, unresolved masks, unsupported sampling filters, and partial/non-addressable host traces fail closed. The adapter is opt-in and does not register or import the vLLM path. Test Plan: - pytest tests/post_training/unit/diffusion_rl_test.py - included in the 83-test focused diffusion-RL suite - Pyink and changed-file Pylint 10.00/10 - py_compile and git diff --check --- src/maxtext/integration/tunix/diffusion_rl.py | 459 +++++++++++++++++ tests/post_training/unit/diffusion_rl_test.py | 470 ++++++++++++++++++ 2 files changed, 929 insertions(+) create mode 100644 src/maxtext/integration/tunix/diffusion_rl.py create mode 100644 tests/post_training/unit/diffusion_rl_test.py diff --git a/src/maxtext/integration/tunix/diffusion_rl.py b/src/maxtext/integration/tunix/diffusion_rl.py new file mode 100644 index 0000000000..bfefcfe6bc --- /dev/null +++ b/src/maxtext/integration/tunix/diffusion_rl.py @@ -0,0 +1,459 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Trajectory-aware MaxText adapters for block-diffusion reinforcement learning.""" + +from collections.abc import Sequence + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.diffusion import scoring +from maxtext.diffusion import denoise +from tunix.diffusion import types as diffusion_types +from tunix.rl import reshard as rl_reshard +from tunix.rl.rollout import base_rollout + + +def _concrete_numpy(value): + if isinstance(value, jax.core.Tracer): + return None + if isinstance(value, jax.Array) and not value.is_fully_addressable: + return None + return np.asarray(value) + + +def prepare_diffusion_policy_batch( + *, + prompt_tokens, + prompt_mask, + completion_tokens, + completion_mask, + action_steps, + loss_mask=None, + inactive_target_id: int = 0, +) -> diffusion_types.DiffusionTokenBatch: + """Builds one completion-aligned policy batch from a sampled denoising trace. + + ``completion_mask`` describes the full canvas used during sampling, while + ``loss_mask`` selects actions owned by the RL objective. They intentionally + differ after an EOS token: exact replay still needs the post-EOS trace as + context for earlier bidirectional actions, but those later actions carry no + policy loss. + """ + expected_shape = tuple(completion_tokens.shape) + if tuple(completion_mask.shape) != expected_shape or tuple(action_steps.shape) != expected_shape: + raise ValueError("completion_tokens, completion_mask, and action_steps must have identical shapes") + if tuple(prompt_tokens.shape) != tuple(prompt_mask.shape): + raise ValueError("prompt_tokens and prompt_mask must have identical shapes") + if prompt_tokens.shape[0] != completion_tokens.shape[0]: + raise ValueError("prompt and completion batches must have the same leading dimension") + if loss_mask is None: + loss_mask = completion_mask + if tuple(loss_mask.shape) != expected_shape: + raise ValueError("loss_mask must match completion_tokens shape") + concrete_mask = _concrete_numpy(completion_mask) + concrete_steps = _concrete_numpy(action_steps) + concrete_loss_mask = _concrete_numpy(loss_mask) + if concrete_mask is not None and concrete_steps is not None and concrete_loss_mask is not None: + active = np.asarray(concrete_mask, dtype=bool) + weighted = np.asarray(concrete_loss_mask, dtype=bool) + concrete_steps = np.asarray(concrete_steps) + if not np.array_equal(concrete_steps >= 0, active): + raise ValueError("each active completion token must have exactly one nonnegative action step") + if np.any(weighted & ~active): + raise ValueError("loss_mask must be a subset of the sampled completion canvas") + if np.any(concrete_steps[active] >= completion_tokens.shape[1]): + raise ValueError("active action steps must be smaller than the completion length") + use_numpy = all( + isinstance(value, np.ndarray) + for value in (prompt_tokens, prompt_mask, completion_tokens, completion_mask, action_steps, loss_mask) + ) + array_module = np if use_numpy else jnp + completion_tokens = array_module.asarray(completion_tokens) + loss_mask = array_module.asarray(loss_mask, dtype=array_module.bool_) + target_ids = array_module.where( + loss_mask, + completion_tokens, + array_module.asarray(inactive_target_id, dtype=completion_tokens.dtype), + ) + return diffusion_types.DiffusionTokenBatch.create( + model_inputs={ + "prompt_tokens": array_module.asarray(prompt_tokens), + "prompt_mask": array_module.asarray(prompt_mask, dtype=array_module.bool_), + "completion_tokens": completion_tokens, + "completion_mask": array_module.asarray(completion_mask, dtype=array_module.bool_), + "action_steps": array_module.asarray(action_steps, dtype=array_module.int32), + }, + target_ids=target_ids, + loss_weights=array_module.asarray(loss_mask, dtype=array_module.float32), + ) + + +def _logical_positions(validity_mask): + positions = jnp.cumsum(jnp.asarray(validity_mask, dtype=jnp.int32), axis=1) - 1 + return jnp.where(validity_mask, positions, 0) + + +def make_diffusion_trace_logits_fn(config): + """Builds a scorer that replays the exact pre-commit canvas for each action.""" + mask_id = int(config.block_diffusion_mask_id) + alignment = config.block_diffusion_logit_alignment + vocab_size = int(config.vocab_size) + + def logits_fn(model: nnx.Module, model_inputs: diffusion_types.ModelInputs): + prompt_tokens = jnp.asarray(model_inputs["prompt_tokens"]) + prompt_mask = jnp.asarray(model_inputs["prompt_mask"], dtype=jnp.bool_) + completion_tokens = jnp.asarray(model_inputs["completion_tokens"]) + completion_mask = jnp.asarray(model_inputs["completion_mask"], dtype=jnp.bool_) + action_steps = jnp.asarray(model_inputs["action_steps"], dtype=jnp.int32) + batch_size, completion_length = completion_tokens.shape + prompt_length = prompt_tokens.shape[1] + validity_mask = jnp.concatenate([prompt_mask, completion_mask], axis=1) + positions = _logical_positions(validity_mask) + segment_ids = validity_mask.astype(jnp.int32) + base_model = getattr(model, "base", model) + graphdef, params, rest = nnx.split(base_model, nnx.Param, ...) + output_logits = jnp.zeros((batch_size, completion_length, vocab_size), dtype=jnp.float32) + + def score_action_step(step, accumulated_logits): + active_actions = completion_mask & (action_steps == step) + + def run_forward(current_logits): + pre_commit_completion = jnp.where(action_steps >= step, mask_id, completion_tokens) + canvas = jnp.concatenate([prompt_tokens, pre_commit_completion], axis=1) + local_model = nnx.merge(graphdef, params, rest, copy=True) + raw_logits = local_model( + decoder_input_tokens=canvas, + decoder_positions=positions, + decoder_segment_ids=segment_ids, + enable_dropout=False, + decoder_target_tokens=canvas, + decoder_target_mask=segment_ids, + ) + target_aligned = scoring.align_logits_to_targets(raw_logits, alignment, positions, validity_mask) + completion_logits = jnp.asarray(target_aligned[:, prompt_length:, :], dtype=jnp.float32) + completion_logits = completion_logits.at[..., mask_id].set(-jnp.inf) + return jnp.where(active_actions[..., None], completion_logits, current_logits) + + return jax.lax.cond(jnp.any(active_actions), run_forward, lambda value: value, accumulated_logits) + + output_logits = jax.lax.fori_loop(0, completion_length, score_action_step, output_logits) + return jnp.where(completion_mask[..., None], output_logits, 0.0) + + return logits_fn + + +def make_diffusion_rollout_fn(config, temperature: float, *, sample_tokens: bool = True): + """Builds one compiled stochastic rollout over the current MaxText policy.""" + max_denoise_steps = config.rl.diffusion_max_denoise_steps + if max_denoise_steps == -1: + max_denoise_steps = config.block_diffusion_block_size + + @nnx.jit + def rollout(model, initial_tokens, positions, segment_ids, completion_mask, rng): + graphdef, params, rest = nnx.split(model, nnx.Param, ...) + validity_mask = segment_ids != 0 + + def canvas_logits(canvas): + local_model = nnx.merge(graphdef, params, rest, copy=True) + base_model = getattr(local_model, "base", local_model) + raw_logits = base_model( # pylint: disable=not-callable + decoder_input_tokens=canvas, + decoder_positions=positions, + decoder_segment_ids=segment_ids, + enable_dropout=False, + decoder_target_tokens=canvas, + decoder_target_mask=segment_ids, + ) + return scoring.align_logits_to_targets( + raw_logits, + config.block_diffusion_logit_alignment, + positions, + validity_mask, + ) + + return denoise.low_confidence_rollout( + canvas_logits, + initial_tokens, + positions, + validity_mask, + completion_mask, + block_size=config.block_diffusion_block_size, + mask_id=config.block_diffusion_mask_id, + logit_alignment=config.block_diffusion_logit_alignment, + canvas_policy=config.block_diffusion_canvas_policy, + confidence_threshold=config.rl.diffusion_confidence_threshold, + temperature=temperature, + max_denoise_steps=max_denoise_steps, + rng=rng if sample_tokens else None, + ) + + return rollout + + +def resolve_stop_token_ids(config, tokenizer_eos_id: int | None) -> tuple[int, ...]: + """Resolves and validates the model's generated completion terminators.""" + configured = tuple(int(token_id) for token_id in config.rl.diffusion_stop_token_ids) + if configured: + stop_token_ids = configured + elif tokenizer_eos_id is not None: + stop_token_ids = (int(tokenizer_eos_id),) + else: + raise ValueError("block-diffusion RL requires an EOS token or explicit diffusion_stop_token_ids") + if len(set(stop_token_ids)) != len(stop_token_ids) or any( + token_id < 0 or token_id >= config.vocab_size for token_id in stop_token_ids + ): + raise ValueError("diffusion stop-token IDs must be unique and inside the model vocabulary") + if config.block_diffusion_mask_id in stop_token_ids: + raise ValueError("the diffusion mask token cannot also be a generated stop token") + return stop_token_ids + + +def policy_mask_at_stop( + tokens: jax.Array, + action_steps: jax.Array, + stop_token_ids: Sequence[int], +) -> jax.Array: + """Selects sampled actions through the first generated stop token.""" + active = action_steps >= 0 + stop_positions = active & jnp.isin(tokens, jnp.asarray(stop_token_ids, tokens.dtype)) + seen_stop = jax.lax.associative_scan(jnp.logical_or, stop_positions, axis=1) + after_stop = jnp.concatenate([jnp.zeros_like(seen_stop[:, :1]), seen_stop[:, :-1]], axis=1) + return active & ~after_stop + + +class MaxTextDiffusionRollout(base_rollout.BaseRollout): + """Correctness-first in-process rollout that exports exact denoising traces.""" + + def __init__( + self, + rollout_actor, + tokenizer, + mesh, + rollout_config, + maxtext_config, + cache_config_or_size=None, + ): + del cache_config_or_size + self._model = rollout_actor + self._tokenizer = tokenizer + self._mesh = mesh + self._rollout_config = rollout_config + self._config = maxtext_config + self._pad_id = tokenizer.pad_id() + if self._pad_id is None: + raise ValueError("block-diffusion RL requires a tokenizer pad token") + self._pad_id = int(self._pad_id) + if not 0 <= self._pad_id < maxtext_config.vocab_size: + raise ValueError("the tokenizer pad token must be inside the model vocabulary") + self._eos_id = tokenizer.eos_id() + self._stop_token_ids = resolve_stop_token_ids(maxtext_config, self._eos_id) + self._rollout_fns = { + (float(rollout_config.temperature), True): make_diffusion_rollout_fn( + maxtext_config, float(rollout_config.temperature) + ) + } + configured_seed = getattr(rollout_config, "seed", None) + if configured_seed is None: + self._base_rng = jax.random.PRNGKey(int(getattr(maxtext_config, "data_shuffle_seed", 0))) + else: + configured_seed = jnp.asarray(configured_seed) + if not configured_seed.shape: + self._base_rng = jax.random.PRNGKey(int(configured_seed)) + elif tuple(configured_seed.shape) == (2,): + self._base_rng = configured_seed + else: + raise ValueError("RolloutConfig.seed must be an integer or one JAX PRNG key") + training_stream = self._sampling_stream_key(rollout_config) + self._rng_streams = {training_stream: self._base_rng} + self._generation_context = None + self._generation_context_invocation = 0 + self._generation_call = 0 + + @property + def mesh(self): + return self._mesh + + def _encode_left_padded(self, prompts, max_prompt_length): + rows = [] + masks = [] + for prompt in prompts: + token_ids = list(self._tokenizer.encode(prompt))[-max_prompt_length:] + padding = max_prompt_length - len(token_ids) + rows.append([self._pad_id] * padding + token_ids) + masks.append([False] * padding + [True] * len(token_ids)) + return np.asarray(rows, dtype=np.int32), np.asarray(masks, dtype=bool) + + def _rollout_fn_for(self, rollout_config): + """Returns a cached rollout compiled for the requested sampling mode.""" + top_k = getattr(rollout_config, "top_k", None) + top_p = getattr(rollout_config, "top_p", 1.0) + if top_p not in (None, 1.0) or top_k not in (None, -1, 1): + raise ValueError("block-diffusion rollout supports categorical sampling or greedy top_k=1") + sample_tokens = top_k != 1 + key = (float(rollout_config.temperature), sample_tokens) + if key not in self._rollout_fns: + self._rollout_fns[key] = make_diffusion_rollout_fn( + self._config, + float(rollout_config.temperature), + sample_tokens=sample_tokens, + ) + return self._rollout_fns[key] + + def _sampling_stream_key(self, rollout_config): + return ( + float(rollout_config.temperature), + getattr(rollout_config, "top_k", None), + getattr(rollout_config, "top_p", 1.0), + bool(getattr(rollout_config, "return_logprobs", False)), + ) + + def _next_rng(self, rollout_config): + """Advances the deterministic RNG stream for one generation call.""" + if self._generation_context is not None: + global_step, mode = self._generation_context + mode_id = 0 if mode == "train" else 1 + rollout_key = jax.random.fold_in(self._base_rng, global_step) + rollout_key = jax.random.fold_in(rollout_key, mode_id) + rollout_key = jax.random.fold_in(rollout_key, self._generation_context_invocation) + rollout_key = jax.random.fold_in(rollout_key, self._generation_call) + self._generation_call += 1 + return rollout_key + stream = self._sampling_stream_key(rollout_config) + if stream not in self._rng_streams: + self._rng_streams[stream] = jax.random.fold_in(self._base_rng, len(self._rng_streams)) + self._rng_streams[stream], rollout_key = jax.random.split(self._rng_streams[stream]) + return rollout_key + + def set_generation_context(self, *, global_step: int, mode): + """Derives resumable rollout keys without coupling train and eval calls.""" + context = (int(global_step), str(mode)) + if context == self._generation_context: + self._generation_context_invocation += 1 + else: + self._generation_context = context + self._generation_context_invocation = 0 + self._generation_call = 0 + + def generate(self, prompts, rollout_config, **kwargs): + del kwargs + rollout_fn = self._rollout_fn_for(rollout_config) + prompt_tokens, prompt_mask = self._encode_left_padded(prompts, rollout_config.max_prompt_length) + batch_size = prompt_tokens.shape[0] + completion_length = rollout_config.max_tokens_to_generate + completion_mask = np.ones((batch_size, completion_length), dtype=bool) + validity_mask = np.concatenate([prompt_mask, completion_mask], axis=1) + positions = _logical_positions(jnp.asarray(validity_mask)) + segment_ids = jnp.asarray(validity_mask, dtype=jnp.int32) + initial_tokens = jnp.concatenate( + [ + jnp.asarray(prompt_tokens), + jnp.full( + (batch_size, completion_length), + self._config.block_diffusion_mask_id, + dtype=jnp.int32, + ), + ], + axis=1, + ) + full_completion_mask = jnp.concatenate( + [ + jnp.zeros(prompt_tokens.shape, dtype=jnp.bool_), + jnp.ones((batch_size, completion_length), dtype=jnp.bool_), + ], + axis=1, + ) + rollout_key = self._next_rng(rollout_config) + trace = rollout_fn( + self._model, + initial_tokens, + positions, + segment_ids, + full_completion_mask, + rollout_key, + ) + completion_tokens = trace.tokens[:, -completion_length:] + completion_steps = trace.action_steps[:, -completion_length:] + completion_logps = trace.action_logps[:, -completion_length:] + output_mask = policy_mask_at_stop( + completion_tokens, + completion_steps, + self._stop_token_ids, + ) + for name, value in ( + ("completion tokens", completion_tokens), + ("action steps", completion_steps), + ("action log probabilities", completion_logps), + ("output mask", output_mask), + ): + if isinstance(value, jax.Array) and not value.is_fully_addressable: + raise RuntimeError( + f"block-diffusion RL requires fully addressable {name}; use the supported Pathways single-controller mode" + ) + host_completion_tokens = np.asarray(jax.device_get(completion_tokens)) + host_completion_steps = np.asarray(jax.device_get(completion_steps)) + host_completion_logps = np.asarray(jax.device_get(completion_logps)) + host_completion_mask = host_completion_steps >= 0 + host_output_mask = np.asarray(jax.device_get(output_mask)) + if not np.all(host_completion_mask) or np.any(host_completion_tokens == self._config.block_diffusion_mask_id): + raise RuntimeError("block-diffusion rollout ended with unresolved mask tokens") + if not np.all(np.isfinite(host_completion_logps[host_completion_mask])): + raise RuntimeError("block-diffusion rollout produced non-finite action log probabilities") + diffusion_batch = prepare_diffusion_policy_batch( + prompt_tokens=prompt_tokens, + prompt_mask=prompt_mask, + completion_tokens=host_completion_tokens, + completion_mask=host_completion_mask, + action_steps=host_completion_steps, + inactive_target_id=self._pad_id, + ) + tokens_out = [] + text_out = [] + logps_out = [] + for tokens_row, active_row, logps_row in zip(host_completion_tokens, host_output_mask, host_completion_logps): + length = int(active_row.sum()) + trimmed_tokens = np.asarray(tokens_row[:length], dtype=np.int32) + tokens_out.append(trimmed_tokens) + text_out.append(self._tokenizer.decode(trimmed_tokens.tolist())) + logps_out.append(np.asarray(logps_row, dtype=np.float32)) + return base_rollout.RolloutOutput( + text=text_out, + logits=None, + tokens=tokens_out, + left_padded_prompt_tokens=prompt_tokens, + logprobs=logps_out if rollout_config.return_logprobs else None, + diffusion_batch=diffusion_batch, + ) + + def update_params(self, params, filter_types=None): + destination = nnx.state(self._model, filter_types) if filter_types is not None else nnx.state(self._model) + nnx.update(self._model, rl_reshard.reshard_pytree(params, destination)) + + def get_per_token_logps(self, prompt_tokens, completion_tokens): + del prompt_tokens, completion_tokens + raise NotImplementedError( + "block-diffusion policy scores require RolloutOutput.diffusion_batch; final-sequence AR scoring is invalid" + ) + + def pad_id(self): + return self._pad_id + + def eos_id(self): + return self._eos_id + + def model(self): + return self._model diff --git a/tests/post_training/unit/diffusion_rl_test.py b/tests/post_training/unit/diffusion_rl_test.py new file mode 100644 index 0000000000..62ce3b7abe --- /dev/null +++ b/tests/post_training/unit/diffusion_rl_test.py @@ -0,0 +1,470 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests exact MaxText denoising-trace replay for diffusion RL.""" + +from types import SimpleNamespace + +from absl.testing import absltest +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.diffusion import denoise +from maxtext.diffusion import scoring +from maxtext.integration.tunix import diffusion_rl +from tunix.rl import diffusion as tunix_diffusion + + +def _config(**overrides): + """Builds the minimal diffusion configuration used by these tests.""" + values = { + "block_diffusion_block_size": 4, + "block_diffusion_mask_id": 7, + "block_diffusion_logit_alignment": "same_position", + "block_diffusion_canvas_policy": "all_masked", + "vocab_size": 8, + "data_shuffle_seed": 3, + "rl": SimpleNamespace( + diffusion_stop_token_ids=[], + diffusion_max_denoise_steps=-1, + diffusion_confidence_threshold=0.99, + ), + } + values.update(overrides) + return SimpleNamespace(**values) + + +class _CanvasDependentModel(nnx.Module): + + def __init__(self): + self.scale = nnx.Param(jnp.asarray(2.0, dtype=jnp.float32)) + + def __call__(self, decoder_input_tokens, **kwargs): + del kwargs + committed = jnp.sum(decoder_input_tokens != 7, axis=1, keepdims=True) + target_ids = (committed + jnp.arange(decoder_input_tokens.shape[1])[None, :]) % 7 + return jax.nn.one_hot(target_ids, 8, dtype=jnp.float32) * self.scale[...] + + +class _NonFiniteModel(nnx.Module): + + def __init__(self): + self.scale = nnx.Param(jnp.asarray(1.0, dtype=jnp.float32)) + + def __call__(self, decoder_input_tokens, **kwargs): + del kwargs + return jnp.full((*decoder_input_tokens.shape, 8), jnp.nan) * self.scale[...] + + +class _Tokenizer: + """Minimal tokenizer contract used by the in-process rollout tests.""" + + def encode(self, text): + del text + return [4] + + def decode(self, token_ids): + return " ".join(str(token_id) for token_id in token_ids) + + def pad_id(self): + return 0 + + def eos_id(self): + return 6 + + +class _PadValuedPromptTokenizer(_Tokenizer): + + def encode(self, text): + del text + return [0, 4] + + +class DiffusionRLTest(absltest.TestCase): + + def test_trace_replay_matches_action_time_logps_and_has_live_gradient(self): + model = _CanvasDependentModel() + initial = jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion_mask = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + temperature = 0.7 + + trace = denoise.low_confidence_rollout( + lambda canvas: model(decoder_input_tokens=canvas), + initial, + positions, + jnp.ones_like(completion_mask), + completion_mask, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + temperature=temperature, + rng=jax.random.PRNGKey(7), + ) + batch = diffusion_rl.prepare_diffusion_policy_batch( + prompt_tokens=initial[:, :1], + prompt_mask=jnp.ones((1, 1), dtype=jnp.bool_), + completion_tokens=trace.tokens[:, 1:], + completion_mask=trace.action_steps[:, 1:] >= 0, + action_steps=trace.action_steps[:, 1:], + ) + logits_fn = diffusion_rl.make_diffusion_trace_logits_fn(_config()) + replayed = tunix_diffusion.compute_diffusion_per_token_logps( + model, + batch, + logits_fn, + temperature=temperature, + stop_gradient=False, + ) + + np.testing.assert_allclose(replayed, trace.action_logps[:, 1:], rtol=1e-6, atol=1e-6) + gradients = nnx.grad( + lambda scored_model: jnp.sum( + tunix_diffusion.compute_diffusion_per_token_logps( + scored_model, + batch, + logits_fn, + temperature=temperature, + stop_gradient=False, + ) + ) + )(model) + self.assertGreater(float(jnp.abs(gradients.scale[...])), 0.0) + + def test_shifted_seed_trace_replay_matches_action_time_logps(self): + model = _CanvasDependentModel() + initial = jnp.asarray([[2, 3, 4, 5, 7, 7, 7, 7]], dtype=jnp.int32) + positions = jnp.arange(8, dtype=jnp.int32)[None, :] + completion_mask = jnp.asarray([[0, 0, 0, 0, 1, 1, 1, 1]], dtype=jnp.bool_) + temperature = 0.7 + + trace = denoise.low_confidence_rollout( + lambda canvas: scoring.align_logits_to_targets( + model(decoder_input_tokens=canvas), + "shifted", + positions, + jnp.ones_like(completion_mask), + ), + initial, + positions, + jnp.ones_like(completion_mask), + completion_mask, + block_size=4, + mask_id=7, + logit_alignment="shifted", + canvas_policy="seed_and_mask", + confidence_threshold=0.99, + temperature=temperature, + rng=jax.random.PRNGKey(9), + ) + batch = diffusion_rl.prepare_diffusion_policy_batch( + prompt_tokens=initial[:, :4], + prompt_mask=jnp.ones((1, 4), dtype=jnp.bool_), + completion_tokens=trace.tokens[:, 4:], + completion_mask=jnp.ones((1, 4), dtype=jnp.bool_), + action_steps=trace.action_steps[:, 4:], + ) + config = _config( + block_diffusion_logit_alignment="shifted", + block_diffusion_canvas_policy="seed_and_mask", + ) + + replayed = tunix_diffusion.compute_diffusion_per_token_logps( + model, + batch, + diffusion_rl.make_diffusion_trace_logits_fn(config), + temperature=temperature, + stop_gradient=True, + ) + + np.testing.assert_allclose(replayed, trace.action_logps[:, 4:], rtol=1e-6, atol=1e-6) + + def test_prepared_batch_rejects_action_mask_mismatch(self): + with self.assertRaisesRegex(ValueError, "exactly one nonnegative action step"): + diffusion_rl.prepare_diffusion_policy_batch( + prompt_tokens=jnp.asarray([[4]]), + prompt_mask=jnp.asarray([[1]], dtype=jnp.bool_), + completion_tokens=jnp.asarray([[5, 6]]), + completion_mask=jnp.asarray([[1, 1]], dtype=jnp.bool_), + action_steps=jnp.asarray([[0, -1]]), + ) + + def test_stop_tokens_mask_actions_after_first_match(self): + keep = diffusion_rl.policy_mask_at_stop( + jnp.asarray([[4, 5, 6, 9]]), + jnp.asarray([[0, 1, 2, 3]]), + (5, 8), + ) + + np.testing.assert_array_equal(keep, [[1, 1, 0, 0]]) + + def test_stop_mask_preserves_full_trace_for_exact_replay(self): + model = _CanvasDependentModel() + initial = jnp.asarray([[4, 7, 7, 7]], dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion_mask = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + temperature = 0.7 + trace = denoise.low_confidence_rollout( + lambda canvas: model(decoder_input_tokens=canvas), + initial, + positions, + jnp.ones_like(completion_mask), + completion_mask, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + temperature=temperature, + rng=jax.random.PRNGKey(7), + ) + completion_tokens = trace.tokens[:, 1:] + completion_steps = trace.action_steps[:, 1:] + stop_id = int(completion_tokens[0, 0]) + policy_mask = diffusion_rl.policy_mask_at_stop(completion_tokens, completion_steps, (stop_id,)) + batch = diffusion_rl.prepare_diffusion_policy_batch( + prompt_tokens=initial[:, :1], + prompt_mask=jnp.ones((1, 1), dtype=jnp.bool_), + completion_tokens=completion_tokens, + completion_mask=jnp.ones_like(completion_tokens, dtype=jnp.bool_), + action_steps=completion_steps, + loss_mask=policy_mask, + inactive_target_id=0, + ) + + replayed = tunix_diffusion.compute_diffusion_per_token_logps( + model, + batch, + diffusion_rl.make_diffusion_trace_logits_fn(_config()), + temperature=temperature, + stop_gradient=True, + ) + + np.testing.assert_array_equal(batch.model_inputs["action_steps"], completion_steps) + np.testing.assert_array_equal(batch.target_ids, jnp.where(policy_mask, completion_tokens, 0)) + np.testing.assert_allclose( + replayed, + jnp.where(policy_mask, trace.action_logps[:, 1:], 0.0), + rtol=1e-6, + atol=1e-6, + ) + + def test_explicit_stop_ids_do_not_require_tokenizer_eos(self): + config = _config( + rl=SimpleNamespace( + diffusion_stop_token_ids=[5, 6], + diffusion_max_denoise_steps=-1, + diffusion_confidence_threshold=0.99, + ) + ) + + self.assertEqual(diffusion_rl.resolve_stop_token_ids(config, None), (5, 6)) + + def test_prompt_mask_comes_from_padding_length_not_token_value(self): + rollout_config = SimpleNamespace(temperature=0.7) + rollout = diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_CanvasDependentModel(), + tokenizer=_PadValuedPromptTokenizer(), + mesh=None, + rollout_config=rollout_config, + maxtext_config=_config(), + ) + + tokens, mask = rollout._encode_left_padded(["prompt"], 3) # pylint: disable=protected-access + + np.testing.assert_array_equal(tokens, [[0, 0, 4]]) + np.testing.assert_array_equal(mask, [[0, 1, 1]]) + + def test_rollout_exports_padded_trace_and_matching_sampler_logps(self): + config = _config() + model = _CanvasDependentModel() + rollout_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.7, + return_logprobs=True, + top_k=-1, + top_p=1.0, + ) + rollout = diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=model, + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=rollout_config, + maxtext_config=config, + ) + + output = rollout.generate(["prompt", "prompt"], rollout_config) + batch = output.diffusion_batch + self.assertIsNotNone(batch) + self.assertIsInstance(batch.target_ids, np.ndarray) + self.assertEqual(batch.target_ids.shape, (2, 3)) + np.testing.assert_array_equal(batch.loss_weights > 0, batch.model_inputs["action_steps"] >= 0) + np.testing.assert_array_equal(batch.target_ids, batch.model_inputs["completion_tokens"]) + replayed = tunix_diffusion.compute_diffusion_per_token_logps( + model, + batch, + diffusion_rl.make_diffusion_trace_logits_fn(config), + temperature=rollout_config.temperature, + stop_gradient=True, + ) + np.testing.assert_allclose(replayed, np.stack(output.logprobs), rtol=1e-6, atol=1e-6) + + def test_rollout_scores_full_trace_while_reward_text_stops_at_eos(self): + config = _config( + rl=SimpleNamespace( + diffusion_stop_token_ids=list(range(7)), + diffusion_max_denoise_steps=-1, + diffusion_confidence_threshold=0.99, + ) + ) + training_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.7, + return_logprobs=True, + top_k=-1, + top_p=1.0, + ) + rollout = diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_CanvasDependentModel(), + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=training_config, + maxtext_config=config, + ) + + output = rollout.generate(["prompt"], training_config) + + self.assertLen(output.tokens[0], 1) + self.assertLen(output.logprobs[0], 3) + np.testing.assert_array_equal(output.diffusion_batch.loss_weights, np.ones((1, 3))) + + def test_greedy_eval_uses_its_own_sampling_contract(self): + training_config = SimpleNamespace(temperature=0.7) + rollout = diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_CanvasDependentModel(), + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=training_config, + maxtext_config=_config(), + ) + eval_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.01, + return_logprobs=False, + top_k=1, + top_p=1.0, + ) + + first = rollout.generate(["prompt"], eval_config) + second = rollout.generate(["prompt"], eval_config) + + np.testing.assert_array_equal(first.diffusion_batch.target_ids, second.diffusion_batch.target_ids) + + def test_rollout_fails_loudly_on_non_finite_policy_scores(self): + rollout_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.7, + return_logprobs=True, + top_k=-1, + top_p=1.0, + ) + rollout = diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_NonFiniteModel(), + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=rollout_config, + maxtext_config=_config(), + ) + + with self.assertRaisesRegex(RuntimeError, "unresolved mask|non-finite"): + rollout.generate(["prompt"], rollout_config) + + def test_eval_rng_does_not_advance_training_stream(self): + training_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.7, + return_logprobs=True, + top_k=-1, + top_p=1.0, + seed=13, + ) + eval_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.01, + return_logprobs=False, + top_k=1, + top_p=1.0, + ) + + def make_rollout(): + return diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_CanvasDependentModel(), + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=training_config, + maxtext_config=_config(), + ) + + with_eval = make_rollout() + baseline = make_rollout() + with_eval.generate(["prompt"], eval_config) + + actual = with_eval.generate(["prompt"], training_config) + expected = baseline.generate(["prompt"], training_config) + + np.testing.assert_array_equal(actual.diffusion_batch.target_ids, expected.diffusion_batch.target_ids) + + def test_generation_context_restores_training_rng_by_global_step(self): + training_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.7, + return_logprobs=True, + top_k=-1, + top_p=1.0, + seed=13, + ) + + def make_rollout(): + return diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_CanvasDependentModel(), + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=training_config, + maxtext_config=_config(), + ) + + uninterrupted = make_rollout() + resumed = make_rollout() + uninterrupted.set_generation_context(global_step=7, mode="train") + resumed.set_generation_context(global_step=7, mode="train") + + expected = uninterrupted.generate(["prompt"], training_config) + actual = resumed.generate(["prompt"], training_config) + + np.testing.assert_array_equal(actual.diffusion_batch.target_ids, expected.diffusion_batch.target_ids) + + +if __name__ == "__main__": + absltest.main()