From 60591b3d629c8da8e97433b21fa6f576ac782363 Mon Sep 17 00:00:00 2001 From: ethannnnnn Date: Wed, 22 Jul 2026 22:29:00 -0700 Subject: [PATCH 1/3] [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/3] [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/3] [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, + )