diff --git a/scripts/ci/pytorch_ci_test_runner.sh b/scripts/ci/pytorch_ci_test_runner.sh index 4ccb827ac1..a8cce192a9 100755 --- a/scripts/ci/pytorch_ci_test_runner.sh +++ b/scripts/ci/pytorch_ci_test_runner.sh @@ -41,7 +41,7 @@ case "$COMMAND" in model_tests) python -m tests.integration_tests.run_tests \ --test_suite models \ - --exclude "qwen3_5_moe_fsdp+tp+ep+pp_spmd_types" \ + --exclude "qwen3_5_moe_fsdp+tp+ep+pp" \ --ngpu "$NGPU" \ "$OUTPUT_DIR" ;; diff --git a/tests/integration_tests/features.py b/tests/integration_tests/features.py index 366c51328b..8bfb7fcc6c 100755 --- a/tests/integration_tests/features.py +++ b/tests/integration_tests/features.py @@ -53,22 +53,22 @@ def _supports_spmd_typechecking(test_name: str, variant: tuple[str, ...]) -> boo """List of tests/variants to test spmd_types backend, but without typechecking.""" unsupported_tests = [ # Compile is not compatible with SPMD typechecking yet. - "1d_compile_spmd_types", - "1d_compile_sac_op_spmd_types", - "2d_compile_spmd_types", - "2d_asynctp_compile_spmd_types", - "3d_compile_spmd_types", - "torchcomms_3d_dp+cp+pp+compile_spmd_types", - "torchcomms_3d_dp+tp+pp+compile_spmd_types", + "1d_compile", + "1d_compile_sac_op", + "2d_compile", + "2d_asynctp_compile", + "3d_compile", + "torchcomms_3d_dp+cp+pp+compile", + "torchcomms_3d_dp+tp+pp+compile", # PP is not compatible with SPMD typechecking yet. - "pp_dp_1f1b_spmd_types", - "pp_tp_gpipe_spmd_types", - "pp_dp_tp_spmd_types", - "validation_tp_cp_pp_spmd_types", - "float8_emulate_lora_spmd_types", + "pp_dp_1f1b", + "pp_tp_gpipe", + "pp_dp_tp", + "validation_tp_cp_pp", + "float8_emulate_lora", # non-chunked CE loss isn't happy yet. ( - "2d_eager_spmd_types", + "2d_eager", [ "--module llama3 --config llama3_debugmodel_ce_loss", "--parallelism.tensor_parallel_degree 2", @@ -81,16 +81,10 @@ def _supports_spmd_typechecking(test_name: str, variant: tuple[str, ...]) -> boo ) -def _enable_spmd_backend(t: OverrideDefinitions, backend: str) -> OverrideDefinitions: - """Inject ``--parallelism.spmd_backend {backend}`` into every variant and - suffix the test name with ``backend``. - - PP-only variants (see ``_is_pp_only``) and CP + compile variants (upstream - symint limitation) skip the backend override. For ``spmd_types``, - ``--debug.spmd_typechecking`` is enabled where supported. Plain (no-backend) - coverage lives in models.py. - """ - test_name = f"{t.test_name}_{backend}" +def _configure_spmd_backend_and_typechecking( + t: OverrideDefinitions, +) -> OverrideDefinitions: + """Configure the spmd_types backend and enable typechecking where supported.""" new_args = [] for variant in t.override_args: prefix: list[str] = [] @@ -98,23 +92,18 @@ def _enable_spmd_backend(t: OverrideDefinitions, backend: str) -> OverrideDefini has_cp = any("context_parallel_degree" in arg for arg in variant) has_compile = any("compile.enable" in arg for arg in variant) has_ac_mode = any("activation-checkpoint:" in arg for arg in variant) - if not _is_pp_only(variant, t.ngpu) and not (has_cp and has_compile): - prefix.append(f"--parallelism.spmd_backend {backend}") - if ( - backend == "spmd_types" - and not has_ac_mode - and _supports_spmd_typechecking(test_name, variant) - ): - prefix.append("--debug.spmd_typechecking") - suffix.append("activation-checkpoint:none") - variant = tuple( - arg.replace(f"{t.test_name}/", f"{test_name}/") for arg in variant - ) + if ( + not _is_pp_only(variant, t.ngpu) + and not (has_cp and has_compile) + and not has_ac_mode + and _supports_spmd_typechecking(t.test_name, variant) + ): + prefix.append("--debug.spmd_typechecking") + suffix.append("activation-checkpoint:none") new_args.append(tuple(prefix) + tuple(variant) + tuple(suffix)) return dataclasses.replace( t, override_args=tuple(new_args), - test_name=test_name, ) @@ -692,5 +681,8 @@ def build_features_test_list() -> list[OverrideDefinitions]: ] return [ - *[_enable_spmd_backend(t, "spmd_types") for t in integration_tests_flavors], + *[ + _configure_spmd_backend_and_typechecking(t) + for t in integration_tests_flavors + ], ] diff --git a/tests/integration_tests/h100.py b/tests/integration_tests/h100.py index d9b281173b..f82f0e1584 100755 --- a/tests/integration_tests/h100.py +++ b/tests/integration_tests/h100.py @@ -103,9 +103,6 @@ def build_h100_tests_list() -> list[OverrideDefinitions]: "--module deepseek_v3 --config " "deepseek_v3_debugmodel_minimal_async_ep", "--compile.no-enable", - # TODO: Remove this per-test override once the H100 suite - # is migrated to the spmd_types backend. - "--parallelism.spmd_backend spmd_types", "--parallelism.data_parallel_shard_degree 2", "--parallelism.context_parallel_degree 2", "--parallelism.tensor_parallel_degree 2", diff --git a/tests/integration_tests/models.py b/tests/integration_tests/models.py index 09011075a5..d0eafdf62f 100755 --- a/tests/integration_tests/models.py +++ b/tests/integration_tests/models.py @@ -10,18 +10,17 @@ from tests.integration_tests import OverrideDefinitions -def _enable_spmd_backend(t: OverrideDefinitions, backend: str) -> OverrideDefinitions: - test_name = f"{t.test_name}_{backend}" +def _configure_spmd_backend_and_typecheck( + t: OverrideDefinitions, +) -> OverrideDefinitions: + """Configure the SPMD backend and enable typechecking where supported.""" + # Compile, PP, and explicit AC modes are not compatible with SPMD + # typechecking yet; keep those as backend-only coverage. new_args = [] for variant in t.override_args: - variant = tuple( - arg.replace(f"{t.test_name}/", f"{test_name}/") for arg in variant - ) - prefix = [f"--parallelism.spmd_backend {backend}"] + prefix = [] suffix = [] - # Compile, PP, and explicit AC modes are not compatible with SPMD - # typechecking yet; keep those as backend-only coverage. - if backend == "spmd_types" and not any( + if not any( token in arg for arg in variant for token in ( @@ -37,7 +36,6 @@ def _enable_spmd_backend(t: OverrideDefinitions, backend: str) -> OverrideDefini return dataclasses.replace( t, override_args=tuple(new_args), - test_name=test_name, ) @@ -265,17 +263,14 @@ def build_model_tests_list() -> list[OverrideDefinitions]: # FSDP+TP+EP+PP test into one supported FSDP+EP path. Kimi # DistMuon rejects TP because it produces _StridedShard # storage. PP support follows in the next stack change. - # Do not enable --debug.spmd_typechecking: multimodal pixel - # tensors from the dataloader are not SPMD-annotated yet. "--module kimi_k2_7 --config kimi_k2_5_debugmodel", - "--parallelism.spmd_backend spmd_types", "--parallelism.data_parallel_shard_degree 4", "--parallelism.expert_parallel_degree 2", "--training.steps 1", ], ], - "Kimi K2.7 DistMuon spmd_types FSDP+EP", - "kimi_k2_5_muon_fsdp+ep_spmd_types", + "Kimi K2.7 DistMuon FSDP+EP", + "kimi_k2_5_muon_fsdp+ep", ngpu=4, ), # Integration Test Cases for Muse Glimmer @@ -294,4 +289,4 @@ def build_model_tests_list() -> list[OverrideDefinitions]: ), ] - return [_enable_spmd_backend(t, "spmd_types") for t in model_tests] + return [_configure_spmd_backend_and_typecheck(t) for t in model_tests] diff --git a/tests/unit_tests/test_dist_gemm.py b/tests/unit_tests/test_dist_gemm.py index 82cb5da69a..f9981cde2f 100644 --- a/tests/unit_tests/test_dist_gemm.py +++ b/tests/unit_tests/test_dist_gemm.py @@ -36,7 +36,7 @@ ) from torchtitan.distributed.parallel_dims import ParallelDims -from torchtitan.distributed.utils import set_spmd_backend +from torchtitan.distributed.utils import get_spmd_backend, set_spmd_backend from torchtitan.models.common.config_utils import make_gqa_config from torchtitan.models.common.decoder_sharding import set_gqa_attention_sharding from torchtitan.models.common.dist_gemm import ( @@ -50,13 +50,14 @@ @contextlib.contextmanager -def spmd_types_backend(): - """dist-GEMM only serves this backend; the sharding setup enforces it.""" - set_spmd_backend("spmd_types") +def use_spmd_backend(backend: str): + """Temporarily select an SPMD backend without leaking test state.""" + previous_backend = get_spmd_backend() + set_spmd_backend(backend) try: yield finally: - set_spmd_backend("default") + set_spmd_backend(previous_backend) class TestDistGemmAttentionConfig(unittest.TestCase): @@ -121,9 +122,11 @@ def test_dtensor_backend_is_rejected(self): attn = model_registry("debugmodel", tp_gemm_backend="dist_gemm") attn = attn.model.layers[0].attention - # the default backend is active here, which is exactly what must be refused - with self.assertRaisesRegex(ValueError, "requires parallelism.spmd_backend"): - set_gqa_attention_sharding(attn, enable_sp=True) + with use_spmd_backend("partial_dtensor"): + with self.assertRaisesRegex( + ValueError, "requires parallelism.spmd_backend" + ): + set_gqa_attention_sharding(attn, enable_sp=True) def test_sequence_parallel_disabled_is_rejected(self): """The fused GEMMs *are* the SP collectives, so SP off has nothing to fuse @@ -132,7 +135,7 @@ def test_sequence_parallel_disabled_is_rejected(self): attn = model_registry("debugmodel", tp_gemm_backend="dist_gemm") attn = attn.model.layers[0].attention - with spmd_types_backend(): + with use_spmd_backend("spmd_types"): with self.assertRaisesRegex(ValueError, "enable_sequence_parallel"): set_gqa_attention_sharding(attn, enable_sp=False) @@ -170,7 +173,7 @@ def test_sharding_setup_declares_the_fused_contracts(self): .model.layers[0] .attention ) - with spmd_types_backend(): + with use_spmd_backend("spmd_types"): set_gqa_attention_sharding(stock, enable_sp=True) set_gqa_attention_sharding(fused, enable_sp=True) @@ -225,11 +228,10 @@ def test_parallelize_preserves_the_fused_contracts(self): parallel_dims = self._parallel_dims() attn_cfg = llama3_debugmodel_dist_gemm().model_spec.model.layers[0].attention - with spmd_types_backend(): + with use_spmd_backend("spmd_types"): set_gqa_attention_sharding(attn_cfg, enable_sp=True) - - attn = attn_cfg.build().to(self.device_type) - attn.parallelize(parallel_dims) + attn = attn_cfg.build().to(self.device_type) + attn.parallelize(parallel_dims) self.assertIsNone(attn._sharding_config) self.assertIsNone(attn.wo._sharding_config.out_src_shardings) @@ -253,7 +255,6 @@ def world_size(self) -> int: @with_comms def test_matches_stock_feed_forward(self): from torchtitan.distributed.spmd_types import set_current_spmd_mesh - from torchtitan.distributed.utils import set_spmd_backend from torchtitan.models.common.config_utils import make_ffn_config R = self.world_size @@ -304,13 +305,10 @@ def test_matches_stock_feed_forward(self): # needs mesh_dim_names, and a "tp" axis for _tp_group_from_context mesh = init_device_mesh(self.device_type, (R,), mesh_dim_names=("tp",)) - set_spmd_backend("spmd_types") - try: + with use_spmd_backend("spmd_types"): with set_current_spmd_mesh(mesh): x_shard = x.chunk(R, 1)[self.rank].contiguous() out_shard = fused(x_shard) - finally: - set_spmd_backend("default") # fused returns this rank's sequence shard of the full-sequence result torch.testing.assert_close( diff --git a/tests/unit_tests/test_embedding.py b/tests/unit_tests/test_embedding.py index 2beb8f07ee..edd05a92ec 100644 --- a/tests/unit_tests/test_embedding.py +++ b/tests/unit_tests/test_embedding.py @@ -183,7 +183,7 @@ def test_vocab_parallel_embedding_parity(self): dst=out_type, ) finally: - set_spmd_backend("default") + set_spmd_backend("spmd_types") # local matches DTensor bitwise and no-parallel embedding self.assertTrue( diff --git a/tests/unit_tests/test_loss.py b/tests/unit_tests/test_loss.py index ea9f56bc4d..e70817c421 100644 --- a/tests/unit_tests/test_loss.py +++ b/tests/unit_tests/test_loss.py @@ -784,7 +784,7 @@ def init_pg(self, eager_init, backend=None): def destroy_pg(self, device_id=None): super().destroy_pg(device_id) - set_spmd_backend("default") + set_spmd_backend("spmd_types") @property def world_size(self): diff --git a/tests/unit_tests/test_module.py b/tests/unit_tests/test_module.py index 490dcc8cad..976a021edb 100644 --- a/tests/unit_tests/test_module.py +++ b/tests/unit_tests/test_module.py @@ -369,6 +369,7 @@ def _parallel_dims(self): pp=1, ep=1, world_size=self.world_size, + spmd_backend="partial_dtensor", ) with patch( "torchtitan.distributed.parallel_dims.device_type", self.device_type diff --git a/tests/unit_tests/test_parallel_dims.py b/tests/unit_tests/test_parallel_dims.py index 575599c574..3649d30c46 100644 --- a/tests/unit_tests/test_parallel_dims.py +++ b/tests/unit_tests/test_parallel_dims.py @@ -506,6 +506,7 @@ def test_single_rank_mesh_operations(self): pp=1, ep=1, world_size=1, + spmd_backend="partial_dtensor", ) # Test mesh building @@ -573,6 +574,7 @@ def test_get_mesh_with_list_input(self): pp=1, ep=1, world_size=1, + spmd_backend="partial_dtensor", ) parallel_dims.build_mesh() @@ -611,8 +613,8 @@ def test_expert_parallelism_validation(self): self.assertTrue(parallel_dims.dp_shard_enabled) -class TestSpmdMeshesLegacy(DTensorTestBase): - """spmd_meshes() under non-full_dtensor.""" +class TestSpmdMeshesPartialDTensor(DTensorTestBase): + """spmd_meshes() under partial_dtensor.""" @property def world_size(self): @@ -631,7 +633,7 @@ def test_legacy_spmd_meshes(self): pp=1, ep=1, world_size=8, - spmd_backend="default", + spmd_backend="partial_dtensor", ) pd.build_mesh() @@ -699,6 +701,7 @@ def test_efsdp_excluded_when_ep_disabled(self): pp=1, ep=1, world_size=8, + spmd_backend="partial_dtensor", ) pd.build_mesh() @@ -769,6 +772,7 @@ def test_world_size_8_mesh_operations(self): pp=1, ep=1, world_size=8, + spmd_backend="partial_dtensor", ) # Test mesh building diff --git a/tests/unit_tests/test_quantization.py b/tests/unit_tests/test_quantization.py index 464cfd2bbc..8d9f3e95d2 100644 --- a/tests/unit_tests/test_quantization.py +++ b/tests/unit_tests/test_quantization.py @@ -242,9 +242,9 @@ def test_nvfp4_recipes_default_to_spmd_types_and_allow_cli_override( assert config.parallelism.spmd_backend == "spmd_types" overridden = ConfigManager().parse_args( - [*base_args, "--parallelism.spmd_backend", "default"] + [*base_args, "--parallelism.spmd_backend", "partial_dtensor"] ) - assert overridden.parallelism.spmd_backend == "default" + assert overridden.parallelism.spmd_backend == "partial_dtensor" @pytest.mark.parametrize( diff --git a/torchtitan/components/loss.py b/torchtitan/components/loss.py index a14c8b38ad..8c5d1bb660 100644 --- a/torchtitan/components/loss.py +++ b/torchtitan/components/loss.py @@ -40,7 +40,7 @@ def cross_entropy_loss( return _cross_entropy_via_local_map(pred, labels) if isinstance(pred, DTensor): - assert get_spmd_backend() == "default" + assert get_spmd_backend() == "partial_dtensor" if pred.placements == (Shard(pred.ndim - 1),): return _LossParallelCrossEntropy.apply( pred.to_local().flatten(0, 1).float(), @@ -327,9 +327,10 @@ def __call__( loss = self.fn(pred, labels) if global_valid_tokens is not None: # TODO(pianpwk): Teach spmd_types that P / scalar preserves P. + is_type_checking = spmd.is_type_checking() with spmd.no_typecheck(): loss = loss / global_valid_tokens - if get_spmd_backend() == "spmd_types": + if is_type_checking: spmd.assert_type(loss, {"dp": spmd.P, "cp": spmd.P, "tp": spmd.I}) return loss, {} @@ -358,9 +359,10 @@ def __call__( loss = self.fn(pred, labels, global_vocab_size=self.global_vocab_size) if global_valid_tokens is not None: # TODO(pianpwk): Teach spmd_types that P / scalar preserves P. + is_type_checking = spmd.is_type_checking() with spmd.no_typecheck(): loss = loss / global_valid_tokens - if get_spmd_backend() == "spmd_types": + if is_type_checking: spmd.assert_type(loss, {"dp": spmd.P, "cp": spmd.P, "tp": spmd.I}) return loss, {} diff --git a/torchtitan/config/configs.py b/torchtitan/config/configs.py index 086afc2c13..8bff05aa47 100644 --- a/torchtitan/config/configs.py +++ b/torchtitan/config/configs.py @@ -160,11 +160,13 @@ class ParallelismConfig: enable_sequence_parallel: bool = True """Whether to use SequenceParallel as part of tensor parallelism. Enabled by default.""" - spmd_backend: Literal["default", "full_dtensor", "spmd_types"] = "default" + spmd_backend: Literal[ + "partial_dtensor", "full_dtensor", "spmd_types" + ] = "spmd_types" """ SPMD backend selector. - - "default": use the existing TorchTitan parallelism paths. + - "partial_dtensor": use DTensor for model-parallel axes only. - "full_dtensor": use the existing full DTensor path. - "spmd_types": use the spmd_types path. """ @@ -248,10 +250,14 @@ class ParallelismConfig: """ def __post_init__(self): - if self.spmd_backend not in {"default", "full_dtensor", "spmd_types"}: + if self.spmd_backend not in { + "partial_dtensor", + "full_dtensor", + "spmd_types", + }: raise ValueError( "parallelism.spmd_backend must be one of " - "'default', 'full_dtensor', or 'spmd_types'." + "'partial_dtensor', 'full_dtensor', or 'spmd_types'." ) if self.context_parallel_load_balancer == "": raise ValueError( diff --git a/torchtitan/distributed/parallel_dims.py b/torchtitan/distributed/parallel_dims.py index a7f724f0ba..bd50347953 100644 --- a/torchtitan/distributed/parallel_dims.py +++ b/torchtitan/distributed/parallel_dims.py @@ -145,7 +145,9 @@ class ParallelDims: pp: int ep: int world_size: int - spmd_backend: Literal["default", "full_dtensor", "spmd_types"] = "default" + spmd_backend: Literal[ + "partial_dtensor", "full_dtensor", "spmd_types" + ] = "spmd_types" # Cache by axis name(s); DeviceMesh equality is by identity, so reuse # is required for ``mesh in spmd_meshes()`` checks. _single_axis_meshes: dict[str, DeviceMesh] = field(default_factory=dict) @@ -578,7 +580,7 @@ def resolve_mesh(self, axes: Iterable[MeshAxisName | str]) -> DeviceMesh | None: axes_list = [ axis.value if isinstance(axis, MeshAxisName) else axis for axis in axes ] - if self.spmd_backend == "default": + if self.spmd_backend == "partial_dtensor": in_band = ("tp", "ep") axes_list = [axis for axis in axes_list if axis in in_band] elif self.spmd_backend == "full_dtensor": diff --git a/torchtitan/distributed/utils.py b/torchtitan/distributed/utils.py index 5efc0d38c6..b392d7540b 100644 --- a/torchtitan/distributed/utils.py +++ b/torchtitan/distributed/utils.py @@ -33,7 +33,7 @@ from torchtitan.distributed.parallel_dims import ParallelDims -_spmd_backend = "default" +_spmd_backend = "spmd_types" def set_spmd_backend(spmd_backend: str) -> None: diff --git a/torchtitan/experiments/graph_trainer/configs.py b/torchtitan/experiments/graph_trainer/configs.py index 4d89e8b808..987b963ab4 100644 --- a/torchtitan/experiments/graph_trainer/configs.py +++ b/torchtitan/experiments/graph_trainer/configs.py @@ -241,6 +241,10 @@ def to_graph_trainer_config( from .trainer import GraphTrainer d = {f.name: getattr(base_config, f.name) for f in fields(base_config)} + d["parallelism"] = replace( + base_config.parallelism, + spmd_backend="partial_dtensor", + ) graph_spec = model_registry(base_config.model_spec.flavor) # Wrap the base model config in the graph_trainer's model config class # (e.g. GraphTrainerQwen3Model.Config) while preserving all field values diff --git a/torchtitan/experiments/graph_trainer/precompile_main.py b/torchtitan/experiments/graph_trainer/precompile_main.py index 41c3b0a048..2cfe7f9271 100644 --- a/torchtitan/experiments/graph_trainer/precompile_main.py +++ b/torchtitan/experiments/graph_trainer/precompile_main.py @@ -102,6 +102,7 @@ def _common_setup(config): pp=pp, ep=parallelism.expert_parallel_degree, world_size=world_size, + spmd_backend=parallelism.spmd_backend, ) parallel_dims.build_mesh() diff --git a/torchtitan/experiments/graph_trainer/tests/_trainer_test_utils.py b/torchtitan/experiments/graph_trainer/tests/_trainer_test_utils.py index c397139197..3c06571ca2 100644 --- a/torchtitan/experiments/graph_trainer/tests/_trainer_test_utils.py +++ b/torchtitan/experiments/graph_trainer/tests/_trainer_test_utils.py @@ -84,14 +84,14 @@ def build_minimal_trainer( parallelism=SimpleNamespace( pipeline_parallel_degree=1, fsdp_reshard_after_forward=fsdp_reshard_after_forward, - spmd_backend="default", + spmd_backend="partial_dtensor", ), ) trainer._fwd_bwd_step_module = None trainer._traced_step = None else: trainer.config = SimpleNamespace( - parallelism=SimpleNamespace(spmd_backend="default"), + parallelism=SimpleNamespace(spmd_backend="partial_dtensor"), ) return trainer diff --git a/torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py b/torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py index 2c5ba52a42..44a465ee54 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py +++ b/torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py @@ -425,6 +425,7 @@ def _setup(self) -> None: pp=1, ep=1, world_size=self.world_size, + spmd_backend="partial_dtensor", ) diff --git a/torchtitan/experiments/graph_trainer/tests/test_numerics.py b/torchtitan/experiments/graph_trainer/tests/test_numerics.py index 7cd6ab5e6c..16f7617c2b 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_numerics.py +++ b/torchtitan/experiments/graph_trainer/tests/test_numerics.py @@ -691,6 +691,7 @@ def init_test(self): pp=1, ep=1, world_size=self.world_size, + spmd_backend="partial_dtensor", ) def get_input(self): diff --git a/torchtitan/experiments/graph_trainer/tests/test_passes.py b/torchtitan/experiments/graph_trainer/tests/test_passes.py index e10b00eeca..04e8e35f93 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_passes.py +++ b/torchtitan/experiments/graph_trainer/tests/test_passes.py @@ -253,6 +253,7 @@ def _setup(self): pp=1, ep=1, world_size=self.world_size, + spmd_backend="partial_dtensor", ) def _make_fsdp_model(self, dim=16, n_layers=3): @@ -1215,6 +1216,7 @@ def _setup(self): pp=1, ep=1, world_size=self.world_size, + spmd_backend="partial_dtensor", ) def _get_fsdp_pg_name(self): @@ -1811,6 +1813,7 @@ def _run_and_get_layer_ids(self, fsdp_reshard_after_forward: str): pp=1, ep=1, world_size=self.world_size, + spmd_backend="partial_dtensor", ) model_spec = llama3_model_registry("debugmodel") diff --git a/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py b/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py index c41d063a21..dd9379580c 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py +++ b/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py @@ -48,6 +48,7 @@ def test_param_cast_to_bf16_at_ngpu_1(self): pp=1, ep=1, world_size=1, + spmd_backend="partial_dtensor", ) training = TrainingConfig( mixed_precision_param="bfloat16", diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py index c944bdfc89..958df6411a 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py +++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py @@ -1842,6 +1842,7 @@ def _setup(self): pp=1, ep=1, world_size=self.world_size, + spmd_backend="partial_dtensor", ) def _run_fsdp_model_test( @@ -2197,6 +2198,7 @@ def test_peak_memory_identical_fsdp(self): pp=1, ep=1, world_size=self.world_size, + spmd_backend="partial_dtensor", ) fsdp_mesh = parallel_dims.get_mesh("fsdp") diff --git a/torchtitan/experiments/rl/models/vllm_registry.py b/torchtitan/experiments/rl/models/vllm_registry.py index 6e7164866f..2bbc05c9f7 100644 --- a/torchtitan/experiments/rl/models/vllm_registry.py +++ b/torchtitan/experiments/rl/models/vllm_registry.py @@ -76,7 +76,7 @@ class InferenceParallelismConfig: enable_sequence_parallel: bool = False """Enable dense sequence parallelism across the tensor-parallel axis.""" - spmd_backend: Literal["default", "spmd_types"] = "default" + spmd_backend: Literal["partial_dtensor", "spmd_types"] = "spmd_types" """SPMD backend used by TorchTitan model parallelization in the generator.""" @property diff --git a/torchtitan/experiments/transformers_modeling_backend/config_registry.py b/torchtitan/experiments/transformers_modeling_backend/config_registry.py index 54378b61ea..a3b7a62470 100644 --- a/torchtitan/experiments/transformers_modeling_backend/config_registry.py +++ b/torchtitan/experiments/transformers_modeling_backend/config_registry.py @@ -45,7 +45,10 @@ def transformers_modeling_backend_debugmodel() -> TransformersBackendConfig: ), dataloader=HuggingFaceTextDataLoader.Config(dataset="c4_test"), metrics=MetricsProcessor.Config(log_freq=1), - parallelism=ParallelismConfig(pipeline_parallel_schedule="1F1B"), + parallelism=ParallelismConfig( + pipeline_parallel_schedule="1F1B", + spmd_backend="partial_dtensor", + ), checkpoint=CheckpointManager.Config( interval=10, last_save_model_only=False, @@ -76,7 +79,10 @@ def transformers_modeling_backend_debugmodel_moe() -> TransformersBackendConfig: ), dataloader=HuggingFaceTextDataLoader.Config(dataset="c4_test"), metrics=MetricsProcessor.Config(log_freq=1), - parallelism=ParallelismConfig(pipeline_parallel_schedule="1F1B"), + parallelism=ParallelismConfig( + pipeline_parallel_schedule="1F1B", + spmd_backend="partial_dtensor", + ), checkpoint=CheckpointManager.Config( interval=10, last_save_model_only=False, @@ -105,7 +111,10 @@ def transformers_modeling_backend_full_moe() -> TransformersBackendConfig: ), dataloader=HuggingFaceTextDataLoader.Config(dataset="c4"), metrics=MetricsProcessor.Config(log_freq=10), - parallelism=ParallelismConfig(pipeline_parallel_schedule="1F1B"), + parallelism=ParallelismConfig( + pipeline_parallel_schedule="1F1B", + spmd_backend="partial_dtensor", + ), checkpoint=CheckpointManager.Config( interval=500, last_save_model_only=False, @@ -136,7 +145,10 @@ def transformers_modeling_backend_full() -> TransformersBackendConfig: ), dataloader=HuggingFaceTextDataLoader.Config(dataset="c4"), metrics=MetricsProcessor.Config(log_freq=1), - parallelism=ParallelismConfig(pipeline_parallel_schedule="1F1B"), + parallelism=ParallelismConfig( + pipeline_parallel_schedule="1F1B", + spmd_backend="partial_dtensor", + ), checkpoint=CheckpointManager.Config( interval=10, last_save_model_only=False, @@ -181,6 +193,10 @@ def process_sample(sample): sample_processor=process_sample, ), metrics=MetricsProcessor.Config(log_freq=1), + parallelism=ParallelismConfig( + pipeline_parallel_schedule="1F1B", + spmd_backend="partial_dtensor", + ), checkpoint=CheckpointManager.Config( enable=True, initial_load_in_hf=True, @@ -232,6 +248,10 @@ def process_sample(sample): sample_processor=process_sample, ), metrics=MetricsProcessor.Config(log_freq=1), + parallelism=ParallelismConfig( + pipeline_parallel_schedule="1F1B", + spmd_backend="partial_dtensor", + ), checkpoint=CheckpointManager.Config( interval=10, last_save_model_only=False, diff --git a/torchtitan/experiments/transformers_modeling_backend/parallelize.py b/torchtitan/experiments/transformers_modeling_backend/parallelize.py index a325ccb68b..5df0f260bc 100644 --- a/torchtitan/experiments/transformers_modeling_backend/parallelize.py +++ b/torchtitan/experiments/transformers_modeling_backend/parallelize.py @@ -101,11 +101,12 @@ def parallelize_hf_transformers( ({parallel_dims.tp}) and 2 * CP degree ({parallel_dims.cp}). """ - # Only the "default" sharding backend is wired here. + # Only the partial-DTensor sharding backend is wired here. # TODO: wire spmd_types (next PR) -- see the migration TODO in hf_sharding.py. - if parallel_dims.spmd_backend != "default": + if parallel_dims.spmd_backend != "partial_dtensor": raise NotImplementedError( - f"The HF transformers backend only supports spmd_backend='default' " + f"The HF transformers backend only supports " + f"spmd_backend='partial_dtensor' " f"today; got '{parallel_dims.spmd_backend}'. spmd_types/full_dtensor " "are not yet wired for this backend (FSDP mesh resolution, " "Titan-native embedding, and attention kernels are pending)." diff --git a/torchtitan/experiments/transformers_modeling_backend/tests/integration_tests.py b/torchtitan/experiments/transformers_modeling_backend/tests/integration_tests.py index 2c3ec24e89..c14babc8b3 100644 --- a/torchtitan/experiments/transformers_modeling_backend/tests/integration_tests.py +++ b/torchtitan/experiments/transformers_modeling_backend/tests/integration_tests.py @@ -32,6 +32,7 @@ def build_transformers_modeling_backend_test_list() -> list[OverrideDefinitions] "--parallelism.expert_parallel_degree 2", "--parallelism.context_parallel_degree 2", "--parallelism.context_parallel_load_balancer ptrr", + "--training.disable_cuda_graphs", "--training.steps 2", ], ], @@ -48,6 +49,7 @@ def build_transformers_modeling_backend_test_list() -> list[OverrideDefinitions] "--parallelism.tensor_parallel_degree 2", "--parallelism.pipeline_parallel_degree 2", "--parallelism.pipeline_parallel_schedule 1F1B", + "--training.disable_cuda_graphs", "--training.steps 2", ], ], @@ -65,6 +67,7 @@ def build_transformers_modeling_backend_test_list() -> list[OverrideDefinitions] "--parallelism.pipeline_parallel_degree 2", "--parallelism.pipeline_parallel_schedule 1F1B", "--parallelism.context_parallel_load_balancer ptrr", + "--training.disable_cuda_graphs", "--training.steps 2", ], ], diff --git a/torchtitan/experiments/transformers_modeling_backend/tests/test_moe_parallelism.py b/torchtitan/experiments/transformers_modeling_backend/tests/test_moe_parallelism.py index c9ecadba9c..f90ad2369b 100644 --- a/torchtitan/experiments/transformers_modeling_backend/tests/test_moe_parallelism.py +++ b/torchtitan/experiments/transformers_modeling_backend/tests/test_moe_parallelism.py @@ -173,7 +173,7 @@ class _FakeParallelDims: """Minimal ParallelDims stub for tests that don't use full distributed setup.""" full_dtensor = False - spmd_backend = "default" + spmd_backend = "partial_dtensor" tp_enabled = False ep_enabled = False tp = 1 diff --git a/torchtitan/models/common/attention.py b/torchtitan/models/common/attention.py index ad4e3f2f2d..51d0658fea 100644 --- a/torchtitan/models/common/attention.py +++ b/torchtitan/models/common/attention.py @@ -770,7 +770,7 @@ def local_qkv_head_split(x): # TODO(pianpwk): this should be doable once spmd_types tracks sharding evenness. with spmd.local(): x_ = x.view(bs, seqlen, -1, self.head_dim) - if get_spmd_backend() == "spmd_types": + if spmd.is_type_checking(): spmd.assert_type( x_, spmd.V, spmd.PartitionSpec("dp", "cp", "tp", None) ) @@ -830,7 +830,7 @@ def forward( # pyrefly: ignore[bad-override] qkv = self.wqkv(x) with spmd.local(): # TODO(pianpwk): same QKV:S(2) unflatten case handled by even sharding qkv = qkv.view(bs, seqlen, -1, self.r_dim, self.head_dim) - if get_spmd_backend() == "spmd_types": + if spmd.is_type_checking(): spmd.assert_type( qkv, spmd.V, spmd.PartitionSpec("dp", "cp", "tp", None, None) ) diff --git a/torchtitan/models/deepseek_v3/model.py b/torchtitan/models/deepseek_v3/model.py index 255ec640e4..6d2cf179b8 100644 --- a/torchtitan/models/deepseek_v3/model.py +++ b/torchtitan/models/deepseek_v3/model.py @@ -11,7 +11,6 @@ import torch from torch import nn -from torchtitan.distributed.utils import get_spmd_backend from torchtitan.models.common.attention import ( AttentionMasksType, BaseAttention, @@ -109,7 +108,7 @@ def forward( # TODO(pianpwk): same QKV:S(2) unflatten case handled by even sharding with spmd.local(): q = q.view(bsz, seqlen, -1, self.qk_head_dim) - if get_spmd_backend() == "spmd_types": + if spmd.is_type_checking(): spmd.assert_type( q, {"dp": spmd.S(0), "cp": spmd.S(1), "tp": spmd.S(2)}, @@ -134,7 +133,7 @@ def forward( kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 ) k = torch.cat([k_nope, k_pe.expand(-1, -1, k_nope.size(2), -1)], dim=-1) - if get_spmd_backend() == "spmd_types" and not torch.compiler.is_compiling(): + if spmd.is_type_checking() and not torch.compiler.is_compiling(): for t in [k, v]: spmd.assert_type( t, diff --git a/torchtitan/models/kimi_k2_7/model.py b/torchtitan/models/kimi_k2_7/model.py index 1f1cf09ac4..9d3ed13ae5 100644 --- a/torchtitan/models/kimi_k2_7/model.py +++ b/torchtitan/models/kimi_k2_7/model.py @@ -193,7 +193,7 @@ def forward( # pyrefly: ignore [bad-override] else: x = tokens - if get_spmd_backend() == "spmd_types": + if spmd.is_type_checking(): # The scatter restores a token-aligned tensor, so text-model DP # resumes as global batch sharding after the multimodal region. spmd.assert_type(x, {"dp": spmd.S(0), "tp": spmd.R}) diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index 5330e12b43..6c53df38ee 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/torchtitan/models/qwen3_5/model.py @@ -934,7 +934,7 @@ def forward( # pyrefly: ignore [bad-override] else: x = tokens - if get_spmd_backend() == "spmd_types": + if spmd.is_type_checking(): # The scatter restores a token-aligned tensor, so text-model DP # resumes as global batch sharding after the multimodal region. spmd.assert_type(x, {"dp": spmd.S(0), "tp": spmd.R})