diff --git a/tests/unit_tests/cpu/test_context_parallel_validation.py b/tests/unit_tests/cpu/test_context_parallel_validation.py index 12c5bdd0e2..672d1ae29d 100644 --- a/tests/unit_tests/cpu/test_context_parallel_validation.py +++ b/tests/unit_tests/cpu/test_context_parallel_validation.py @@ -5,11 +5,15 @@ # LICENSE file in the root directory of this source tree. import unittest +from unittest import mock import pytest +import torch +from torch.nn.attention.flex_attention import BlockMask from torchtitan.config import ParallelismConfig -from torchtitan.distributed.context_parallel import validate_cp_backend +from torchtitan.distributed.context_parallel import cp_shard, validate_cp_backend +from torchtitan.distributed.pipeline_parallel import pipeline_vlm class TestValidateCpBackend(unittest.TestCase): @@ -31,6 +35,78 @@ def test_allows_partial_dtensor_without_cp(self): validate_cp_backend(self._parallelism(spmd_backend="partial_dtensor", cp=1)) +class TestContextParallelMaskSharding(unittest.TestCase): + def test_mixed_mask_mapping_preserves_non_block_mask_metadata(self): + input_T = torch.arange(8) + input_shard_T = input_T[:4] + block_mask = mock.Mock(spec=BlockMask) + sharded_block_mask = mock.Mock(spec=BlockMask) + varlen_metadata = mock.sentinel.varlen_metadata + cp_context = mock.sentinel.cp_context + attention_masks = { + "quadratic_attention": block_mask, + "deltanet": varlen_metadata, + "deltanet_cp_context": cp_context, + } + cp_mesh = mock.Mock() + cp_mesh.size.return_value = 2 + + with mock.patch( + "torchtitan.distributed.context_parallel.api._context_parallel_shard", + side_effect=[[input_shard_T], [sharded_block_mask]], + ): + sharded_inputs, sharded_masks = cp_shard( + cp_mesh, + (input_T,), + attention_masks, + load_balancer_type=None, + ) + + self.assertIs(sharded_inputs[0], input_shard_T) + assert isinstance(sharded_masks, dict) + self.assertIs(sharded_masks["quadratic_attention"], sharded_block_mask) + self.assertIs(sharded_masks["deltanet"], varlen_metadata) + self.assertIs(sharded_masks["deltanet_cp_context"], cp_context) + + +class TestVlmPipelineInputModules(unittest.TestCase): + def test_post_scatter_reshard_stays_with_token_embeddings(self): + model = mock.Mock() + model.decoder_input_reshard = mock.Mock() + parallelism = ParallelismConfig( + module_fqns_per_model_part=[ + ["vision_encoder", "tok_embeddings", "layers.0"], + ["layers.1", "norm", "lm_head"], + ] + ) + expected = mock.sentinel.pipeline_result + + with mock.patch( + "torchtitan.distributed.pipeline_parallel.pipeline_llm", + return_value=expected, + ) as pipeline_llm: + result = pipeline_vlm( + model, + parallel_dims=mock.sentinel.parallel_dims, + parallelism=parallelism, + model_config=mock.sentinel.model_config, + ) + + self.assertIs(result, expected) + stage_fqns = pipeline_llm.call_args.kwargs[ + "parallelism" + ].module_fqns_per_model_part + self.assertEqual( + stage_fqns[0], + [ + "vision_encoder", + "tok_embeddings", + "decoder_input_reshard", + "layers.0", + ], + ) + + class TestDecoderConfigCpValidation(unittest.TestCase): """``Decoder.Config.update_from_config`` applies the CP gates at config time.""" @@ -94,5 +170,63 @@ def test_allows_partial_dtensor_without_cp(self): config.model_spec.model.update_from_config(config=config) +class TestQwen35ConfigCpValidation(unittest.TestCase): + @staticmethod + def _config(): + try: + from torchtitan.models.qwen3_5.config_registry import qwen35_debugmodel + except ModuleNotFoundError as exc: + raise unittest.SkipTest( + f"Qwen3.5 optional dependency unavailable: {exc.name}" + ) from exc + + config = qwen35_debugmodel() + config.parallelism.spmd_backend = "spmd_types" + config.parallelism.context_parallel_degree = 2 + config.training.max_context_length = 512 + return config + + def test_rejects_context_parallel_load_balancing(self): + config = self._config() + config.parallelism.context_parallel_load_balancer = "headtail" + with self.assertRaisesRegex(ValueError, "contiguous sequence shards"): + config.model_spec.model.update_from_config( # pyrefly: ignore[missing-attribute] + config=config + ) + + def test_allows_contiguous_context_parallel_sharding(self): + import spmd_types as spmd + + from torchtitan.distributed.parallel_dims import MeshAxisName + from torchtitan.distributed.spmd_types import ( + _per_axis_types, + spmd_validate_redistributions, + ) + + config = self._config() + config.parallelism.context_parallel_load_balancer = None + config.model_spec.model.update_from_config( # pyrefly: ignore[missing-attribute] + config=config + ) + + model_config = config.model_spec.model + reshard_config = model_config.decoder_input_reshard.sharding_config + assert reshard_config is not None + assert reshard_config.in_src_shardings is not None + assert reshard_config.in_dst_shardings is not None + self.assertEqual( + _per_axis_types(reshard_config.in_src_shardings["input"])[MeshAxisName.CP], + spmd.R, + ) + self.assertEqual( + _per_axis_types(reshard_config.in_dst_shardings["input"])[MeshAxisName.CP], + spmd.S(0), + ) + spmd_validate_redistributions(reshard_config) + first_layer_config = model_config.layers[0].sharding_config + assert first_layer_config is not None + spmd_validate_redistributions(first_layer_config) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/cpu/test_parallel_dims.py b/tests/unit_tests/cpu/test_parallel_dims.py index 7ea7f896e1..9943a6ece1 100644 --- a/tests/unit_tests/cpu/test_parallel_dims.py +++ b/tests/unit_tests/cpu/test_parallel_dims.py @@ -275,6 +275,14 @@ def test_decoder_layout_partition_spec_ranks(self): attention_activation_placement().partition_spec, ((MeshAxisName.DP, MeshAxisName.CP), MeshAxisName.TP, None), ) + self.assertEqual( + token_id_placement(cp=spmd.R).partition_spec, + (MeshAxisName.DP,), + ) + self.assertEqual( + _per_axis_types(token_id_placement(cp=spmd.R))[MeshAxisName.CP], + spmd.R, + ) def test_unfold_dp_axes(self): """Logical DP expands only when resolving concrete mesh axes.""" diff --git a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py index 7e2d4ede00..92fd49d83b 100644 --- a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py +++ b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py @@ -144,7 +144,9 @@ def forward( *, cu_seqlens: torch.Tensor | None = None, cu_seqlens_cpu: torch.Tensor | None = None, + cp_context: object | None = None, ) -> torch.Tensor: + assert cp_context is None if xq_THK.shape[1] != xv_THV.shape[1]: assert xv_THV.shape[1] % xq_THK.shape[1] == 0 repeat = xv_THV.shape[1] // xq_THK.shape[1] @@ -168,6 +170,41 @@ def forward( class TestQwen35DeltaNetVarlen(unittest.TestCase): + def test_chunk_kernel_forwards_cp_context(self): + try: + from torchtitan.models.qwen3_5 import GatedDeltaKernel + except ModuleNotFoundError as exc: + raise unittest.SkipTest( + f"Qwen3.5 optional dependency unavailable: {exc.name}" + ) from exc + + kernel = GatedDeltaKernel(GatedDeltaKernel.Config(backend="fla_chunked")) + xq_THK = torch.randn(8, 2, 4) + xk_THK = torch.randn(8, 2, 4) + xv_THV = torch.randn(8, 2, 4) + g_TH = torch.randn(8, 2) + beta_TH = torch.randn(8, 2) + cp_context = mock.sentinel.cp_context + + with mock.patch( + "torchtitan.models.qwen3_5.gdn._fla_chunk_gated_delta_rule", + return_value=(xv_THV.unsqueeze(0), None), + ) as chunk_gated_delta_rule: + output_THV = kernel( + xq_THK, + xk_THK, + xv_THV, + g_TH, + beta_TH, + cp_context=cp_context, + ) + + torch.testing.assert_close(output_THV, xv_THV) + kwargs = chunk_gated_delta_rule.call_args.kwargs + self.assertIs(kwargs["cp_context"], cp_context) + self.assertIsNone(kwargs["cu_seqlens"]) + self.assertIsNone(kwargs["cu_seqlens_cpu"]) + def test_flex_masks_ignore_padding_position_resets(self): try: from torchtitan.models.common.decoder import Decoder diff --git a/tests/unit_tests/test_qwen3_5_mrope_positions.py b/tests/unit_tests/test_qwen3_5_mrope_positions.py index ee59c5c1fd..afdfad255f 100644 --- a/tests/unit_tests/test_qwen3_5_mrope_positions.py +++ b/tests/unit_tests/test_qwen3_5_mrope_positions.py @@ -19,6 +19,7 @@ """ import unittest +from unittest import mock import torch from torch import nn @@ -129,6 +130,85 @@ def test_multimodal_batch_routes_mrope_to_layers(self): batch["attention_masks"]["deltanet"].cu_seq_q_host, (0, 3, 5, 10) ) + def test_context_parallel_context_is_built_before_input_sharding(self): + import spmd_types as spmd + + from torchtitan.distributed.parallel_dims import MeshAxisName + from torchtitan.distributed.spmd_types import _per_axis_types + from torchtitan.models.common.attention import VarlenMetadata + + model, _sink, _parallel_dims, parallelism = self._build_stub_model() + positions = torch.arange(8, dtype=torch.int32) + cu_seqlens = torch.tensor([0, 8], dtype=torch.int32) + deltanet_metadata = VarlenMetadata( + cu_seq_q=cu_seqlens, + cu_seq_k=cu_seqlens, + max_q=8, + max_k=8, + cu_seq_q_host=(0, 8), + ) + attention_masks = { + "quadratic_attention": None, + "deltanet": deltanet_metadata, + } + cp_mesh = mock.Mock() + cp_mesh.size.return_value = 2 + cp_group = mock.sentinel.cp_group + cp_mesh.get_group.return_value = cp_group + parallel_dims = mock.Mock() + parallel_dims.cp_enabled = True + parallel_dims.get_mesh.return_value = cp_mesh + parallelism.context_parallel_load_balancer = None + cp_context = mock.sentinel.cp_context + input_dict = { + "input": torch.randint(0, 100, (8,)), + "positions": positions, + "labels": torch.zeros(8), + "pixel_values": torch.randn(4, 8), + "grid_thw": torch.tensor([[1, 2, 2]]), + } + + with ( + mock.patch.object( + model, + "get_attention_masks", + return_value=attention_masks, + ), + mock.patch( + "torchtitan.models.qwen3_5.model.build_cp_context", + return_value=cp_context, + ) as build_context, + mock.patch( + "torchtitan.distributed.context_parallel.api.prepare_context_parallel_input", + side_effect=lambda batch, *_args: batch, + ) as prepare_cp_input, + ): + _inputs, _labels, batch = model.preprocess_inputs( + input_dict, + parallel_dims=parallel_dims, + parallelism=parallelism, + ) + + self.assertIs(batch["attention_masks"]["deltanet_cp_context"], cp_context) + build_context.assert_called_once() + args, kwargs = build_context.call_args + self.assertIs(args[0], cu_seqlens) + self.assertIs(kwargs["group"], cp_group) + self.assertEqual(kwargs["conv1d_kernel_size"], model.gdn_conv_kernel_size) + torch.testing.assert_close( + kwargs["cu_seqlens_cpu"], torch.tensor([0, 8], dtype=torch.long) + ) + + input_sharding = prepare_cp_input.call_args.args[1] + self.assertEqual( + _per_axis_types(input_sharding["input"])[MeshAxisName.CP], spmd.R + ) + self.assertEqual( + _per_axis_types(input_sharding["pixel_values"])[MeshAxisName.CP], + spmd.R, + ) + self.assertIs(batch["pixel_values"], input_dict["pixel_values"]) + if __name__ == "__main__": unittest.main() diff --git a/torchtitan/distributed/context_parallel/api.py b/torchtitan/distributed/context_parallel/api.py index a70c7c9333..2dfa1e8000 100644 --- a/torchtitan/distributed/context_parallel/api.py +++ b/torchtitan/distributed/context_parallel/api.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from collections.abc import Mapping from typing import Any, cast, TYPE_CHECKING import spmd_types as spmd @@ -71,7 +72,8 @@ def prepare_context_parallel_input( 'labels', and any extra kwargs. Tensor entries named in ``shard_dims`` (e.g. 'input', 'labels', 'positions') are sharded and written back; 'attention_masks', if present, is sharded along its Q - seq dim. + seq dim. In a mixed mapping, non-``BlockMask`` metadata is preserved + unchanged for model-specific CP implementations. input_shardings: Per-input SPMD layout; the CP sequence dim for each input is derived via ``_cp_shard_dims`` (inputs whose CP axis is Replicate/Partial are omitted and left untouched). When None, @@ -141,8 +143,9 @@ def cp_shard( cp_mesh: Device mesh for context parallel dimension inputs: Tuple of input tensors to be sharded along the sequence dimension - attention_masks: Attention masks to be sharded. Supports None, - BlockMask, or dict[str, BlockMask] + attention_masks: Attention masks to be sharded. Supports ``None``, a + ``BlockMask``, or a mapping containing ``BlockMask`` values and + model-specific metadata. Only ``BlockMask`` values are sharded. load_balancer_type: Type of load balancer to use. Options: - "headtail": Use HeadTailLoadBalancer (for SDPA) - "ptrr": Use PTRRLoadBalancer (for FlexAttention) @@ -197,7 +200,7 @@ def cp_shard( "PTRRLoadBalancer requires attention_masks to be a " "BlockMask or dict[str, BlockMask], but got None" ) - if isinstance(attention_masks, dict): + if isinstance(attention_masks, Mapping): if ptrr_mask_key is None: raise ValueError( "PTRRLoadBalancer received a dict[str, BlockMask] " @@ -236,36 +239,35 @@ def cp_shard( ), ) - # BlockMask, has shape, [B, H, Q, KV], and we can only shard - # on the Q seq dimension, not KV. + # BlockMask has shape [B, H, Q, KV], and we can only shard on the Q + # sequence dimension, not KV. Preserve non-BlockMask entries in mixed + # mappings; model-specific linear-attention metadata remains global or + # rank-local according to its own CP implementation. MASK_Q_SEQ_DIM = 2 - if attention_masks is not None: - assert isinstance(attention_masks, (BlockMask, dict)) - masks: list[BlockMask] = [] - for mask in ( - [attention_masks] - if isinstance(attention_masks, BlockMask) - else attention_masks.values() - ): - if not isinstance(mask, BlockMask): - raise ValueError( - "Context parallelism can only shard BlockMask attention " - f"masks, got {type(mask).__name__} in the mask dict." - ) - masks.append(mask) + if isinstance(attention_masks, BlockMask): sharded_masks = _context_parallel_shard( mesh=cp_mesh, - buffers=masks, - seq_dims=(MASK_Q_SEQ_DIM,) * len(masks), + buffers=[attention_masks], + seq_dims=[MASK_Q_SEQ_DIM], load_balancer=load_balancer, ) - attention_masks = cast( - (BlockMask | dict[str, BlockMask]), - ( - sharded_masks[0] - if isinstance(attention_masks, BlockMask) - else {k: v for k, v in zip(attention_masks.keys(), sharded_masks)} - ), + attention_masks = cast(BlockMask, sharded_masks[0]) + elif isinstance(attention_masks, Mapping): + block_mask_items = [ + (key, mask) + for key, mask in attention_masks.items() + if isinstance(mask, BlockMask) + ] + if not block_mask_items: + return inputs, attention_masks + sharded_masks = _context_parallel_shard( + mesh=cp_mesh, + buffers=[mask for _, mask in block_mask_items], + seq_dims=(MASK_Q_SEQ_DIM,) * len(block_mask_items), + load_balancer=load_balancer, ) + attention_masks = dict(attention_masks) + for (key, _), sharded_mask in zip(block_mask_items, sharded_masks, strict=True): + attention_masks[key] = cast(BlockMask, sharded_mask) return inputs, attention_masks diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py index f547b18cd4..6e7e615133 100644 --- a/torchtitan/distributed/pipeline_parallel.py +++ b/torchtitan/distributed/pipeline_parallel.py @@ -156,10 +156,11 @@ def pipeline_vlm( (``tok_embeddings``, ``layers.*``, ``norm``, ``lm_head``). For a VLM we inject ``vision_encoder`` into the first stage's FQN list so it runs alongside ``tok_embeddings`` (vision features are scattered into the embedding sequence - before the decoder layers). On stages other than the first, ``tok_embeddings`` - and ``vision_encoder`` are pruned to ``None``; each model's ``forward`` must - guard on ``self.tok_embeddings is not None`` so the multimodal logic is - skipped there. + before the decoder layers). A model-specific ``decoder_input_reshard`` module, + when present, is also kept with ``tok_embeddings`` so post-scatter CP sharding + runs on the first stage. On later stages these input modules are pruned to + ``None``; each model's ``forward`` must guard on + ``self.tok_embeddings is not None`` so the multimodal logic is skipped there. NOTE: This adds load to stage 0 that the auto split does not model (``input_weight`` only accounts for ``tok_embeddings``); for a heavy vision @@ -178,9 +179,32 @@ def pipeline_vlm( ) if model.vision_encoder is not None: fqn_per_part[0].insert(0, "vision_encoder") - parallelism = dataclasses.replace( - parallelism, module_fqns_per_model_part=fqn_per_part + else: + fqn_per_part = [ + list(stage_fqns) for stage_fqns in parallelism.module_fqns_per_model_part + ] + + if getattr(model, "decoder_input_reshard", None) is not None: + input_stage = next( + ( + stage_fqns + for stage_fqns in fqn_per_part + if "tok_embeddings" in stage_fqns + ), + None, ) + if input_stage is None: + raise ValueError( + "VLM pipeline partition must place tok_embeddings on a stage " + "before decoder_input_reshard can be assigned." + ) + if "decoder_input_reshard" not in input_stage: + embedding_index = input_stage.index("tok_embeddings") + input_stage.insert(embedding_index + 1, "decoder_input_reshard") + + parallelism = dataclasses.replace( + parallelism, module_fqns_per_model_part=fqn_per_part + ) return pipeline_llm( model, diff --git a/torchtitan/models/common/decoder_sharding.py b/torchtitan/models/common/decoder_sharding.py index 028284c1a4..406702c56f 100644 --- a/torchtitan/models/common/decoder_sharding.py +++ b/torchtitan/models/common/decoder_sharding.py @@ -61,8 +61,17 @@ def dense_activation_placement( ) -def token_id_placement() -> SpmdType: +def token_id_placement(*, cp: spmd.PerMeshAxisSpmdType = spmd.S(0)) -> SpmdType: """Placement for decoder token IDs with shape ``(tokens,)``.""" + if not isinstance(cp, spmd.Shard): + return SpmdType( + { + DP: spmd.V, + CP: cp, + TP: spmd.R, + }, + partition_spec=spmd.PartitionSpec(DP), + ) return SpmdType( { DP: spmd.V, @@ -272,7 +281,11 @@ def set_gqa_attention_sharding(attention_cfg, *, enable_sp: bool) -> None: attention_cfg.wo.sharding_config = wo_config -def set_gqa_inner_attention_local_map(inner_attention_cfg) -> None: +def set_gqa_inner_attention_local_map( + inner_attention_cfg, + *, + cp: spmd.PerMeshAxisSpmdType = spmd.S(0), +) -> None: """Install a ``LocalMapConfig`` on an inner-attention config. q/k use ``(T, H, K)`` and v uses ``(T, H, V)``. DP/CP shard T and TP @@ -285,16 +298,18 @@ def set_gqa_inner_attention_local_map(inner_attention_cfg) -> None: is multi-axis); under ``partial_dtensor``, the (tp,)-only mesh only consumes the ``TP`` placement and the rest are ignored. - With CP, q stays token-sharded on the CP axis while k/v are - unsharded (``R``) on CP -- the local_map boundary all-gathers k/v so the - kernel sees full-length keys (matching the BlockMask's kv dimension). - Q's local grad is naturally token-sharded; k/v's local grads accumulate as - partial (``P``) on CP and are reduced on the way out. + For decoder attention, the default keeps q token-sharded on CP while k/v + are unsharded so the kernel sees full-length keys. Vision callers pass + ``cp=R`` because every CP rank processes the complete vision sequence. """ - q_placements = attention_activation_placement() - kv_src_placements = attention_activation_placement() - kv_dst_placements = attention_activation_placement(cp=spmd.R) - kv_grad_placements = attention_activation_placement(cp=spmd.P) + q_placements = attention_activation_placement(cp=cp) + kv_src_placements = attention_activation_placement(cp=cp) + if isinstance(cp, spmd.Shard): + kv_dst_placements = attention_activation_placement(cp=spmd.R) + kv_grad_placements = attention_activation_placement(cp=spmd.P) + else: + kv_dst_placements = kv_src_placements + kv_grad_placements = kv_src_placements out_src: SpmdType = q_placements inner_attention_cfg.sharding_config = ShardingConfig( diff --git a/torchtitan/models/common/multimodal.py b/torchtitan/models/common/multimodal.py index f318eaada2..094eba5ca1 100644 --- a/torchtitan/models/common/multimodal.py +++ b/torchtitan/models/common/multimodal.py @@ -27,8 +27,11 @@ def multimodal_context() -> contextlib.AbstractContextManager[None]: Under ``spmd_types`` the vision encoder and the vision->text scatter run per-DP-rank on that rank's own images: the pixel tensors are DP-local (``V@DP``), so the region must execute with DP treated as a local axis. - After the scatter the tensor is token-aligned again and global DP batch - sharding resumes. A no-op outside ``spmd_types`` (or when DP is size 1). + CP remains active, so a caller may keep the inputs replicated across CP, + finish the scatter on the complete sequence, and shard that sequence at a + later module boundary. After the scatter the tensor is token-aligned again + and global DP batch sharding resumes. A no-op outside ``spmd_types`` (or + when DP is size 1). """ if get_spmd_backend() == "spmd_types" and spmd_mesh_size("dp") > 1: return spmd.set_current_mesh(local_axes=("dp",)) diff --git a/torchtitan/models/common/vision_encoder_sharding.py b/torchtitan/models/common/vision_encoder_sharding.py index 52944a7fe8..3debc99405 100644 --- a/torchtitan/models/common/vision_encoder_sharding.py +++ b/torchtitan/models/common/vision_encoder_sharding.py @@ -20,18 +20,33 @@ DP = MeshAxisName.DP +CP = MeshAxisName.CP TP = MeshAxisName.TP -def multimodal_input_sharding() -> dict[str, SpmdType]: +def _vision_placement( + *, + dp: spmd.PerMeshAxisSpmdType, + tp: spmd.PerMeshAxisSpmdType, + cp: spmd.PerMeshAxisSpmdType | None = None, +) -> SpmdType: + axis_types = {DP: dp} + if cp is not None: + axis_types[CP] = cp + axis_types[TP] = tp + return SpmdType(axis_types) + + +def multimodal_input_sharding( + *, cp: spmd.PerMeshAxisSpmdType | None = None +) -> dict[str, SpmdType]: """SPMD layouts for VLM vision inputs (folded into a model's input_sharding). The vision tensors are DP-local (``V@DP``) -- each DP rank owns its own - images -- and TP-invariant (``I@TP``): the model consumes them inside - ``multimodal_context`` (a DP-local mesh) and the vision encoder runs per-rank. - Shared by every VLM decoder (Qwen3.5, Kimi K2.5, Muse Glimmer). + images -- and TP-invariant (``I@TP``). Callers that prepare a complete + multimodal embedding sequence before CP sharding pass ``cp=R``. """ - layout = SpmdType({DP: spmd.V, TP: spmd.I}) + layout = _vision_placement(dp=spmd.V, cp=cp, tp=spmd.I) return { "pixel_values": layout, "pixel_values_videos": layout, @@ -40,68 +55,76 @@ def multimodal_input_sharding() -> dict[str, SpmdType]: } -def invariant_norm_config() -> ShardingConfig: +def invariant_norm_config( + *, cp: spmd.PerMeshAxisSpmdType | None = None +) -> ShardingConfig: """Norm whose state and activations are invariant across TP ranks.""" return ShardingConfig( state_shardings={ - "weight": SpmdType({DP: spmd.R, TP: spmd.I}), - "bias": SpmdType({DP: spmd.R, TP: spmd.I}), + "weight": _vision_placement(dp=spmd.R, cp=cp, tp=spmd.I), + "bias": _vision_placement(dp=spmd.R, cp=cp, tp=spmd.I), }, in_src_shardings={ - "input": SpmdType({DP: spmd.V, TP: spmd.I}), + "input": _vision_placement(dp=spmd.V, cp=cp, tp=spmd.I), }, in_dst_shardings={ - "input": SpmdType({DP: spmd.V, TP: spmd.I}), + "input": _vision_placement(dp=spmd.V, cp=cp, tp=spmd.I), }, - out_src_shardings=SpmdType({DP: spmd.V, TP: spmd.I}), - out_dst_shardings=SpmdType({DP: spmd.V, TP: spmd.I}), + out_src_shardings=_vision_placement(dp=spmd.V, cp=cp, tp=spmd.I), + out_dst_shardings=_vision_placement(dp=spmd.V, cp=cp, tp=spmd.I), ) -def vision_invariant_linear_config() -> ShardingConfig: +def vision_invariant_linear_config( + *, cp: spmd.PerMeshAxisSpmdType | None = None +) -> ShardingConfig: """Unsharded linear whose state and activations are invariant at TP.""" return ShardingConfig( state_shardings={ - "weight": SpmdType({DP: spmd.R, TP: spmd.I}), - "bias": SpmdType({DP: spmd.R, TP: spmd.I}), + "weight": _vision_placement(dp=spmd.R, cp=cp, tp=spmd.I), + "bias": _vision_placement(dp=spmd.R, cp=cp, tp=spmd.I), }, in_src_shardings={ - "input": SpmdType({DP: spmd.V, TP: spmd.I}), + "input": _vision_placement(dp=spmd.V, cp=cp, tp=spmd.I), }, in_dst_shardings={ - "input": SpmdType({DP: spmd.V, TP: spmd.I}), + "input": _vision_placement(dp=spmd.V, cp=cp, tp=spmd.I), }, - out_src_shardings=SpmdType({DP: spmd.V, TP: spmd.I}), - out_dst_shardings=SpmdType({DP: spmd.V, TP: spmd.I}), + out_src_shardings=_vision_placement(dp=spmd.V, cp=cp, tp=spmd.I), + out_dst_shardings=_vision_placement(dp=spmd.V, cp=cp, tp=spmd.I), ) def vision_colwise_config( - *, input_tp: spmd.PerMeshAxisSpmdType = spmd.I + *, + input_tp: spmd.PerMeshAxisSpmdType = spmd.I, + cp: spmd.PerMeshAxisSpmdType | None = None, ) -> ShardingConfig: """Colwise vision linear with a TP-replicated local matmul input.""" return ShardingConfig( state_shardings={ - "weight": SpmdType({DP: spmd.R, TP: spmd.S(0)}), - "bias": SpmdType({DP: spmd.R, TP: spmd.S(0)}), + "weight": _vision_placement(dp=spmd.R, cp=cp, tp=spmd.S(0)), + "bias": _vision_placement(dp=spmd.R, cp=cp, tp=spmd.S(0)), }, in_src_shardings={ - "input": SpmdType({DP: spmd.V, TP: input_tp}), + "input": _vision_placement(dp=spmd.V, cp=cp, tp=input_tp), }, in_dst_shardings={ - "input": SpmdType({DP: spmd.V, TP: spmd.R}), + "input": _vision_placement(dp=spmd.V, cp=cp, tp=spmd.R), }, - out_src_shardings=SpmdType({DP: spmd.V, TP: spmd.S(-1)}), + out_src_shardings=_vision_placement(dp=spmd.V, cp=cp, tp=spmd.S(-1)), ) -def vision_scaled_bias_rowwise_config() -> ShardingConfig: +def vision_scaled_bias_rowwise_config( + *, cp: spmd.PerMeshAxisSpmdType | None = None +) -> ShardingConfig: """Scaled-bias rowwise vision linear returning a TP-invariant activation.""" - input_layout = SpmdType({DP: spmd.V, TP: spmd.S(1)}) + input_layout = _vision_placement(dp=spmd.V, cp=cp, tp=spmd.S(1)) return ShardingConfig( state_shardings={ - "weight": SpmdType({DP: spmd.R, TP: spmd.S(1)}), - "bias": SpmdType({DP: spmd.R, TP: spmd.R}), + "weight": _vision_placement(dp=spmd.R, cp=cp, tp=spmd.S(1)), + "bias": _vision_placement(dp=spmd.R, cp=cp, tp=spmd.R), }, in_src_shardings={ "input": input_layout, @@ -109,8 +132,8 @@ def vision_scaled_bias_rowwise_config() -> ShardingConfig: in_dst_shardings={ "input": input_layout, }, - out_src_shardings=SpmdType({DP: spmd.V, TP: spmd.P}), - out_dst_shardings=SpmdType({DP: spmd.V, TP: spmd.I}), + out_src_shardings=_vision_placement(dp=spmd.V, cp=cp, tp=spmd.P), + out_dst_shardings=_vision_placement(dp=spmd.V, cp=cp, tp=spmd.I), local_map=LocalMapConfig(in_grad_placements=(input_layout,)), ) @@ -119,26 +142,30 @@ def set_vision_transformer_block_sharding_config( block: "VisionTransformerBlock.Config", *, rope_cache_dp: spmd.PerMeshAxisSpmdType, + cp: spmd.PerMeshAxisSpmdType | None = None, ) -> None: """Set TP sharding for the common vision transformer block.""" - block.norm1.sharding_config = invariant_norm_config() - block.norm2.sharding_config = invariant_norm_config() + block.norm1.sharding_config = invariant_norm_config(cp=cp) + block.norm2.sharding_config = invariant_norm_config(cp=cp) block.attn.sharding_config = ShardingConfig( in_src_shardings={ - "x": SpmdType({DP: spmd.V, TP: spmd.I}), - "rope_cache": SpmdType({DP: rope_cache_dp, TP: spmd.I}), + "x": _vision_placement(dp=spmd.V, cp=cp, tp=spmd.I), + "rope_cache": _vision_placement(dp=rope_cache_dp, cp=cp, tp=spmd.I), }, in_dst_shardings={ - "x": SpmdType({DP: spmd.V, TP: spmd.R}), - "rope_cache": SpmdType({DP: rope_cache_dp, TP: spmd.R}), + "x": _vision_placement(dp=spmd.V, cp=cp, tp=spmd.R), + "rope_cache": _vision_placement(dp=rope_cache_dp, cp=cp, tp=spmd.R), }, ) - block.attn.wq.sharding_config = vision_colwise_config(input_tp=spmd.R) - block.attn.wk.sharding_config = vision_colwise_config(input_tp=spmd.R) - block.attn.wv.sharding_config = vision_colwise_config(input_tp=spmd.R) - block.attn.proj.sharding_config = vision_scaled_bias_rowwise_config() - set_gqa_inner_attention_local_map(block.attn.inner_attention) - - block.mlp.fc1.sharding_config = vision_colwise_config() - block.mlp.fc2.sharding_config = vision_scaled_bias_rowwise_config() + block.attn.wq.sharding_config = vision_colwise_config(input_tp=spmd.R, cp=cp) + block.attn.wk.sharding_config = vision_colwise_config(input_tp=spmd.R, cp=cp) + block.attn.wv.sharding_config = vision_colwise_config(input_tp=spmd.R, cp=cp) + block.attn.proj.sharding_config = vision_scaled_bias_rowwise_config(cp=cp) + if cp is None: + set_gqa_inner_attention_local_map(block.attn.inner_attention) + else: + set_gqa_inner_attention_local_map(block.attn.inner_attention, cp=cp) + + block.mlp.fc1.sharding_config = vision_colwise_config(cp=cp) + block.mlp.fc2.sharding_config = vision_scaled_bias_rowwise_config(cp=cp) diff --git a/torchtitan/models/qwen3_5/README.md b/torchtitan/models/qwen3_5/README.md index 0ab86d5f3c..683cf7832f 100644 --- a/torchtitan/models/qwen3_5/README.md +++ b/torchtitan/models/qwen3_5/README.md @@ -62,6 +62,7 @@ pip install av torchvision flash-linear-attention |---------|-------| | FSDP / HSDP | Decoder sharded per-layer; vision encoder sharded as a single unit (one AllGather) | | Tensor Parallelism (TP) | With Sequence Parallel; head-sharded TP on GatedDeltaNet projections | +| Context Parallelism (CP) | Text and multimodal; vision embeddings are scattered into the complete sequence before decoder CP sharding. GatedDeltaNet uses FLA CP and requires `fla_chunked` with `context_parallel_load_balancer=None` | | Expert Parallelism (EP) | For MoE variants | | Pipeline Parallel (PP) | Vision encoder assigned to first stage; 1F1B and Interleaved1F1B schedules | | Sample Packing | Opt-in via `MMSamplePackingConfig` | @@ -79,4 +80,3 @@ Test scripts: ## TODO - Add video dataset training configs -- Add Context Parallel (CP) support diff --git a/torchtitan/models/qwen3_5/gdn.py b/torchtitan/models/qwen3_5/gdn.py index aa774351dd..07784a2a50 100644 --- a/torchtitan/models/qwen3_5/gdn.py +++ b/torchtitan/models/qwen3_5/gdn.py @@ -12,7 +12,9 @@ import spmd_types as spmd import torch import torch.nn.functional as F +from fla.modules.conv.cp import CausalConv1dFunctionCP from fla.modules.conv.triton.ops import CausalConv1dFunction +from fla.ops.cp import FLACPContext from fla.ops.gated_delta_rule import ( chunk_gated_delta_rule as _fla_chunk_gated_delta_rule, fused_recurrent_gated_delta_rule as _fla_fused_recurrent_gated_delta_rule, @@ -36,6 +38,7 @@ spmd.register_local_autograd_function(ChunkGatedDeltaRuleFunction) spmd.register_local_autograd_function(FusedRecurrentFunction) spmd.register_local_autograd_function(CausalConv1dFunction) +spmd.register_local_autograd_function(CausalConv1dFunctionCP) @spmd.local_map( @@ -78,6 +81,25 @@ def _causal_conv1d_varlen( return out_1TD.squeeze(0) +def _causal_conv1d_cp( + x_TD: torch.Tensor, + weight: torch.Tensor, + cp_context: FLACPContext, +) -> torch.Tensor: + """FLA depthwise causal conv over a context-parallel token shard.""" + from fla.modules.conv.causal_conv1d import causal_conv1d as _fla_causal_conv1d + + out_BTD, _ = _fla_causal_conv1d( + x=x_TD.unsqueeze(0), + weight=weight.squeeze(1), + bias=None, + activation="silu", + backend="triton", + cp_context=cp_context, + ) + return out_BTD.squeeze(0) + + class RMSNormGated(Module): """Gated RMSNorm: ``silu(gate) * weight * norm(x)``. @@ -267,6 +289,7 @@ def forward( *, cu_seqlens: torch.Tensor | None = None, cu_seqlens_cpu: torch.Tensor | None = None, + cp_context: FLACPContext | None = None, ) -> torch.Tensor: # Expand Q/K heads to match V when n_value_heads > n_key_heads if xq_THK.shape[1] != xv_THV.shape[1]: @@ -281,6 +304,18 @@ def forward( g_1TH = g_TH.unsqueeze(0) beta_1TH = beta_TH.unsqueeze(0) + if cp_context is not None and self.backend != "fla_chunked": + raise ValueError( + "Gated DeltaNet context parallelism requires the fla_chunked " + "backend." + ) + + if is_in_batch_invariant_mode() and cp_context is not None: + raise ValueError( + "Gated DeltaNet context parallelism is not supported in " + "batch-invariant recurrent mode." + ) + if is_in_batch_invariant_mode() and cu_seqlens is not None: if cu_seqlens_cpu is None: raise ValueError( @@ -297,7 +332,7 @@ def forward( ).squeeze(0) if self.backend == "fla_chunked": - if cu_seqlens is not None and cu_seqlens_cpu is None: + if cp_context is None and cu_seqlens is not None and cu_seqlens_cpu is None: raise ValueError( "Qwen3.5 FLA varlen DeltaNet requires a CPU cu_seqlens tensor." ) @@ -308,8 +343,9 @@ def forward( g_1TH, beta_1TH, use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, + cu_seqlens=None if cp_context is not None else cu_seqlens, + cu_seqlens_cpu=(None if cp_context is not None else cu_seqlens_cpu), + cp_context=cp_context, ) elif self.backend == "fla_fused_recurrent": result = _fla_fused_recurrent_gated_delta_rule( @@ -364,6 +400,7 @@ def forward( key_head_dim: int, value_head_dim: int, cu_seqlens_host: tuple[int, ...] | None = None, + cp_context: FLACPContext | None = None, ) -> torch.Tensor: """Run separate Q/K/V convolutions and recurrence on local heads.""" num_tokens = query_TC.shape[0] @@ -381,6 +418,9 @@ def causal_conv( x_TC: torch.Tensor, weight_C1W: torch.Tensor, ) -> torch.Tensor: + if cp_context is not None: + return _causal_conv1d_cp(x_TC, weight_C1W, cp_context) + if cu_seqlens_host is not None: return _causal_conv1d_varlen( x_TC, @@ -425,6 +465,7 @@ def causal_conv( beta_TH, cu_seqlens=cu_seqlens if cu_seqlens_host is not None else None, cu_seqlens_cpu=cu_seqlens_cpu, + cp_context=cp_context, ) @@ -489,21 +530,29 @@ def forward( self, x_TD: torch.Tensor, attention_masks: VarlenMetadata | None = None, + *, + cp_context: FLACPContext | None = None, ) -> torch.Tensor: num_tokens = x_TD.shape[0] cu_seqlens_host = None if attention_masks is not None: - # FLA caches varlen index helpers by tensor identity. A fresh - # tensor ensures forward and activation-checkpoint recompute both - # execute the helpers instead of taking different cache paths. - with spmd.local(): - cu_seqlens = attention_masks.cu_seq_q.clone() - cu_seqlens_host = attention_masks.cu_seq_q_host - if cu_seqlens_host is None: - raise ValueError( - "Qwen3.5 GatedDeltaNet varlen requires CPU cu_seqlens " - "metadata. Build VarlenMetadata with include_host_offsets=True." - ) + if cp_context is not None: + # FLA CP carries the rank-local offsets; this replicated tensor + # only preserves the InnerGatedDeltaNet local-map contract. + cu_seqlens = attention_masks.cu_seq_q + else: + # FLA caches varlen index helpers by tensor identity. A fresh + # tensor ensures forward and activation-checkpoint recompute both + # execute the helpers instead of taking different cache paths. + with spmd.local(): + cu_seqlens = attention_masks.cu_seq_q.clone() + cu_seqlens_host = attention_masks.cu_seq_q_host + if cu_seqlens_host is None: + raise ValueError( + "Qwen3.5 GatedDeltaNet varlen requires CPU cu_seqlens " + "metadata. Build VarlenMetadata with " + "include_host_offsets=True." + ) else: cu_seqlens = torch.arange( 0, @@ -537,6 +586,7 @@ def forward( key_head_dim=self.key_head_dim, value_head_dim=self.value_head_dim, cu_seqlens_host=cu_seqlens_host, + cp_context=cp_context, ) gate_THV = gate_TC.view(num_tokens, -1, self.value_head_dim) output_THV = self.norm(output_THV, gate_THV) diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index 7de7c2a84d..e468bc258e 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/torchtitan/models/qwen3_5/model.py @@ -5,11 +5,12 @@ # LICENSE file in the root directory of this source tree. -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, cast import spmd_types as spmd import torch +from fla.ops.cp import build_cp_context, FLACPContext from spmd_types import SpmdType from torch import nn from torch.nn.attention.flex_attention import BlockMask @@ -37,6 +38,7 @@ multimodal_context, scatter_vision_embeds, ) +from torchtitan.models.common.nn_modules import Identity from torchtitan.models.common.vision_encoder_sharding import multimodal_input_sharding from torchtitan.models.utils import ( delta_rule_flops_per_token, @@ -45,7 +47,7 @@ ) from torchtitan.protocols.module import Module -from .gdn import GatedDeltaNet +from .gdn import GatedDeltaKernel, GatedDeltaNet, InnerGatedDeltaNet from .rope import MRoPE from .sharding import annotate_deltanet_cu_seqlens, set_qwen35_sharding_config from .vision_encoder import Qwen35VisionEncoder @@ -56,7 +58,7 @@ # K = query/key head dimension, V = value head dimension, # R = rotary dimension, P = non-rotary dimension. -Qwen35AttentionMaskDict = dict[str, BlockMask | VarlenMetadata | None] +Qwen35AttentionMaskDict = dict[str, BlockMask | VarlenMetadata | FLACPContext | None] class OffsetRMSNorm(Module): @@ -227,14 +229,37 @@ def forward( attention_masks: Qwen35AttentionMaskDict | None, positions: torch.Tensor | None = None, ) -> torch.Tensor: - layer_mask = ( - attention_masks[self.attn_mask_key] if attention_masks is not None else None - ) h_TD = self.attention_norm(x_TD) if self.full_attn: + layer_mask = ( + attention_masks[self.attn_mask_key] + if attention_masks is not None + else None + ) + assert isinstance(layer_mask, (BlockMask, VarlenMetadata)) or ( + layer_mask is None + ) h_TD = self.attn(h_TD, layer_mask, positions) else: - h_TD = self.attn(h_TD, layer_mask) + deltanet_metadata = ( + attention_masks[self.attn_mask_key] + if attention_masks is not None + else None + ) + cp_context = ( + attention_masks.get("deltanet_cp_context") + if attention_masks is not None + else None + ) + assert isinstance(deltanet_metadata, VarlenMetadata) or ( + deltanet_metadata is None + ) + assert isinstance(cp_context, FLACPContext) or cp_context is None + h_TD = self.attn( + h_TD, + deltanet_metadata, + cp_context=cp_context, + ) x_TD = x_TD + h_TD h_TD = self.ffn_norm(x_TD) @@ -281,6 +306,8 @@ class Qwen35Model(Decoder): │ ├─ get_vision_positions → locate vision regions │ └─ _scatter_vision_embeds → scatter into text sequence │ + ├─ decoder_input_reshard CP-shard fused embeddings + │ └─ transformer layers (hybrid), each given ``positions`` (3D or 2D) └─ for each layer: ├─ full attention (every Nth): QK-norm → partial RoPE → SDPA → gate @@ -291,6 +318,7 @@ class Qwen35Model(Decoder): @dataclass(kw_only=True, slots=True) class Config(Decoder.Config): vision_encoder: Qwen35VisionEncoder.Config + decoder_input_reshard: Identity.Config = field(default_factory=Identity.Config) def update_from_config( self, @@ -301,6 +329,29 @@ def update_from_config( Decoder.Config.update_from_config(self, config=config, **kwargs) parallelism = config.parallelism + if ( + parallelism.context_parallel_degree > 1 + and parallelism.context_parallel_load_balancer is not None + ): + raise ValueError( + "Qwen3.5 GatedDeltaNet context parallelism requires " + "context_parallel_load_balancer=None because FLA CP uses " + "contiguous sequence shards." + ) + if parallelism.context_parallel_degree > 1: + for layer in self.layers: + if layer.delta_net is None: + continue + inner_config = layer.delta_net.inner_gated_delta_net + assert isinstance(inner_config, InnerGatedDeltaNet.Config) + kernel_config = inner_config.kernel + assert isinstance(kernel_config, GatedDeltaKernel.Config) + if kernel_config.backend != "fla_chunked": + raise ValueError( + "Qwen3.5 GatedDeltaNet context parallelism requires " + "the fla_chunked backend." + ) + tp = parallelism.tensor_parallel_degree if tp > 1: dn_cfg = next( @@ -365,7 +416,20 @@ def __init__(self, config: Config): super().__init__(config) self.vision_encoder = config.vision_encoder.build() + self.decoder_input_reshard = config.decoder_input_reshard.build() self.spatial_merge_size = config.vision_encoder.spatial_merge_size + deltanet_configs = [ + layer.delta_net for layer in config.layers if layer.delta_net is not None + ] + assert deltanet_configs, "Qwen3.5 must contain GatedDeltaNet layers." + conv_kernel_sizes = { + deltanet_config.conv_kernel_size for deltanet_config in deltanet_configs + } + assert len(conv_kernel_sizes) == 1, ( + "All Qwen3.5 GatedDeltaNet layers must use the same convolution " + "kernel size." + ) + self.gdn_conv_kernel_size = next(iter(conv_kernel_sizes)) def preprocess_inputs( self, @@ -374,7 +438,7 @@ def preprocess_inputs( parallel_dims: ParallelDims, parallelism: ParallelismConfig, ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: - """Build masks, CP-shard, SPMD-wrap (+ deltanet annotation), and return.""" + """Build masks, CP-shard decoder inputs, annotate SPMD types, and return.""" # Function-local import avoids a circular import. from torchtitan.distributed.context_parallel.api import ( prepare_context_parallel_input, @@ -389,7 +453,20 @@ def preprocess_inputs( if isinstance(inner, (FlexAttention.Config, VarlenAttention.Config)): batch["attention_masks"] = self.get_attention_masks(positions=positions) - input_sharding = {**decoder_input_sharding(), **multimodal_input_sharding()} + input_sharding = { + **decoder_input_sharding(), + **multimodal_input_sharding(cp=spmd.R), + } + # Keep token IDs complete through embedding and vision scatter. The + # decoder_input_reshard boundary shards the fused embeddings afterward. + input_sharding["input"] = SpmdType( + { + MeshAxisName.DP: spmd.V, + MeshAxisName.CP: spmd.R, + MeshAxisName.TP: spmd.R, + }, + partition_spec=spmd.PartitionSpec(MeshAxisName.DP), + ) # RoPE uses the 3D MRoPE positions when present (multimodal), else the # same 2D positions. Collapse both into the single ``positions`` input. @@ -415,7 +492,69 @@ def preprocess_inputs( "'positions' or 'mrope_positions'." ) batch["positions"] = rope_positions + attention_masks = batch.get("attention_masks") if parallel_dims.cp_enabled: + if parallelism.context_parallel_load_balancer is not None: + raise ValueError( + "Qwen3.5 GatedDeltaNet context parallelism requires " + "context_parallel_load_balancer=None because FLA CP uses " + "contiguous sequence shards." + ) + if not isinstance(attention_masks, dict): + raise ValueError( + "Qwen3.5 context parallelism requires attention metadata " + "as a keyed mapping." + ) + if positions is None: + raise ValueError( + "Qwen3.5 context parallelism requires 1D text positions." + ) + + deltanet_metadata = attention_masks.get("deltanet") + if deltanet_metadata is None: + num_tokens = positions.shape[0] + cu_seqlens = torch.tensor( + [0, num_tokens], + dtype=torch.int32, + device=positions.device, + ) + deltanet_metadata = VarlenMetadata( + cu_seq_q=cu_seqlens, + cu_seq_k=cu_seqlens, + max_q=num_tokens, + max_k=num_tokens, + cu_seq_q_host=(0, num_tokens), + ) + attention_masks["deltanet"] = deltanet_metadata + if not isinstance(deltanet_metadata, VarlenMetadata): + raise ValueError( + "Qwen3.5 context parallelism requires GatedDeltaNet " + "VarlenMetadata." + ) + if deltanet_metadata.cu_seq_q_host is None: + raise ValueError( + "Qwen3.5 context parallelism requires host cu_seqlens." + ) + + cp_mesh = parallel_dims.get_mesh("cp") + num_tokens = deltanet_metadata.cu_seq_q_host[-1] + if num_tokens % cp_mesh.size() != 0: + raise ValueError( + f"FLA context parallelism requires the token count " + f"({num_tokens}) to be divisible by the CP degree " + f"({cp_mesh.size()})." + ) + cu_seqlens_cpu = torch.tensor( + deltanet_metadata.cu_seq_q_host, + dtype=torch.long, + device="cpu", + ) + attention_masks["deltanet_cp_context"] = build_cp_context( + deltanet_metadata.cu_seq_q, + group=cp_mesh.get_group(), + conv1d_kernel_size=self.gdn_conv_kernel_size, + cu_seqlens_cpu=cu_seqlens_cpu, + ) batch = prepare_context_parallel_input( batch, input_sharding, @@ -592,6 +731,9 @@ def forward( # pyrefly: ignore [bad-override] else: x = tokens + if self.tok_embeddings is not None: + x = self.decoder_input_reshard(x) + if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): spmd.assert_type( x, diff --git a/torchtitan/models/qwen3_5/parallelize.py b/torchtitan/models/qwen3_5/parallelize.py index 1974ddc755..44cbfd5c05 100644 --- a/torchtitan/models/qwen3_5/parallelize.py +++ b/torchtitan/models/qwen3_5/parallelize.py @@ -56,13 +56,6 @@ def parallelize_qwen3_5( compile_config.enable and "model" in compile_config.components ) - if parallel_dims.cp_enabled: - raise NotImplementedError( - "Context Parallel is not yet supported for Qwen3.5. " - "GatedDeltaNet (75% of layers) requires full-sequence allgather, " - "and multimodal CP needs vision scatter before CP sharding." - ) - if ( parallelism.spmd_backend == "spmd_types" or parallel_dims.tp_enabled diff --git a/torchtitan/models/qwen3_5/sharding.py b/torchtitan/models/qwen3_5/sharding.py index bce9c70221..251de8047c 100644 --- a/torchtitan/models/qwen3_5/sharding.py +++ b/torchtitan/models/qwen3_5/sharding.py @@ -74,7 +74,11 @@ def annotate_deltanet_cu_seqlens(attention_masks: "Qwen35AttentionMaskDict") -> return spmd.assert_type( deltanet_metadata.cu_seq_q, - {MeshAxisName.DP: spmd.V, MeshAxisName.TP: spmd.R}, + { + MeshAxisName.DP: spmd.V, + MeshAxisName.CP: spmd.R, + MeshAxisName.TP: spmd.R, + }, ) @@ -120,18 +124,26 @@ def set_qwen35_sharding_config( ) -> None: """Fill ``sharding_config`` on all Qwen3.5 sub-configs.""" set_decoder_sharding_config(config, enable_sp=enable_sp) - # Vision scatter needs the full embedding sequence on every TP rank. + # Vision scatter needs the full embedding sequence on every CP/TP rank. + full_token_layout = token_id_placement(cp=spmd.R) + full_embedding_layout = dense_activation_placement(tp=spmd.R, cp=spmd.R) config.tok_embeddings.sharding_config = ShardingConfig( state_shardings={"weight": dense_param_placement(tp=spmd.S(0))}, - in_src_shardings={"input": token_id_placement()}, - in_dst_shardings={"input": token_id_placement()}, - out_src_shardings=dense_activation_placement(tp=spmd.P, cp=spmd.S(0)), - out_dst_shardings=dense_activation_placement(tp=spmd.R, cp=spmd.S(0)), + in_src_shardings={"input": full_token_layout}, + in_dst_shardings={"input": full_token_layout}, + out_src_shardings=dense_activation_placement(tp=spmd.P, cp=spmd.R), + out_dst_shardings=full_embedding_layout, local_map=LocalMapConfig(in_grad_placements=None), ) + cp_sharded_embedding_layout = dense_activation_placement(tp=spmd.R, cp=spmd.S(0)) + config.decoder_input_reshard.sharding_config = ShardingConfig( + in_src_shardings={"input": full_embedding_layout}, + in_dst_shardings={"input": cp_sharded_embedding_layout}, + out_src_shardings=cp_sharded_embedding_layout, + ) _set_vision_encoder_sharding(config.vision_encoder) - # The first layer restores the decoder layout after replicated vision scatter. - first_layer_input_layout = dense_activation_placement(tp=spmd.R, cp=spmd.S(0)) + # CP sharding happens after vision scatter; layer 0 only restores TP/SP. + first_layer_input_layout = cp_sharded_embedding_layout layer_input_layout = ( dense_sequence_parallel_placement() if enable_sp @@ -247,26 +259,27 @@ def _set_vision_encoder_sharding(ve_cfg: "Qwen35VisionEncoder.Config") -> None: Norms are Replicate. pos_embed is Replicate via state_shardings. """ ve_cfg.sharding_config = ShardingConfig( - state_shardings={"pos_embed": SpmdType({DP: spmd.R, TP: spmd.I})}, - out_src_shardings=SpmdType({DP: spmd.V, TP: spmd.I}), - out_dst_shardings=SpmdType({DP: spmd.V, TP: spmd.R}), + state_shardings={"pos_embed": SpmdType({DP: spmd.R, CP: spmd.R, TP: spmd.I})}, + out_src_shardings=SpmdType({DP: spmd.V, CP: spmd.R, TP: spmd.I}), + out_dst_shardings=SpmdType({DP: spmd.V, CP: spmd.R, TP: spmd.R}), ) ve_cfg.rotary_pos_emb.sharding_config = ShardingConfig( - state_shardings={"inv_freq": SpmdType({DP: spmd.R, TP: spmd.I})}, - out_src_shardings=SpmdType({DP: spmd.R, TP: spmd.I}), + state_shardings={"inv_freq": SpmdType({DP: spmd.R, CP: spmd.R, TP: spmd.I})}, + out_src_shardings=SpmdType({DP: spmd.R, CP: spmd.R, TP: spmd.I}), ) - ve_cfg.patch_embed_proj.sharding_config = vision_invariant_linear_config() + ve_cfg.patch_embed_proj.sharding_config = vision_invariant_linear_config(cp=spmd.R) set_vision_transformer_block_sharding_config( ve_cfg.block, rope_cache_dp=spmd.V, + cp=spmd.R, ) # Merger sub-modules merger = ve_cfg.merger - merger.norm.sharding_config = invariant_norm_config() - merger.fc1.sharding_config = vision_colwise_config() - merger.fc2.sharding_config = vision_scaled_bias_rowwise_config() + merger.norm.sharding_config = invariant_norm_config(cp=spmd.R) + merger.fc1.sharding_config = vision_colwise_config(cp=spmd.R) + merger.fc2.sharding_config = vision_scaled_bias_rowwise_config(cp=spmd.R) def _set_full_attention_sharding( @@ -334,6 +347,13 @@ def _set_deltanet_sharding( projected_placement = dense_activation_placement(tp=spmd.S(1), cp=spmd.S(0)) head_placement = attention_activation_placement() parameter_placement = dense_param_placement(tp=spmd.S(0)) + parameter_grad_placement = SpmdType( + { + DP: spmd.R, + CP: spmd.P, + TP: spmd.S(0), + } + ) replicated_placement = dense_param_placement(tp=spmd.R) cu_seqlens_placement = SpmdType( { @@ -390,19 +410,20 @@ def _set_deltanet_sharding( out_src_shardings=head_placement, out_dst_shardings=head_placement, local_map=LocalMapConfig( - # cu_seqlens varies across DP ranks and is replicated across TP. - # It has no gradient, but local_map still requires its placement. + # Parameter gradients are partial over CP because each rank sees a + # different token shard. cu_seqlens has no gradient, but local_map + # still requires its placement. in_grad_placements=( projected_placement, projected_placement, projected_placement, projected_placement, projected_placement, - parameter_placement, - parameter_placement, - parameter_placement, - parameter_placement, - parameter_placement, + parameter_grad_placement, + parameter_grad_placement, + parameter_grad_placement, + parameter_grad_placement, + parameter_grad_placement, cu_seqlens_placement, ), ),