From c9906f46be8eea8d37ae80e42624c94693bd2f2b Mon Sep 17 00:00:00 2001 From: Pian Pawakapan Date: Fri, 14 Aug 2026 16:40:08 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../experiments/graph_trainer/common_utils.py | 11 ++- .../graph_trainer/dynamic_shapes.py | 4 + .../graph_trainer/llama3/parallelize.py | 2 +- .../graph_trainer/make_fx_tracer.py | 31 +++++--- .../experiments/graph_trainer/simple_fsdp.py | 77 +++++++++++++++++-- .../experiments/graph_trainer/trainer.py | 4 +- 6 files changed, 108 insertions(+), 21 deletions(-) diff --git a/torchtitan/experiments/graph_trainer/common_utils.py b/torchtitan/experiments/graph_trainer/common_utils.py index 11265f6d1c..50f670bb79 100644 --- a/torchtitan/experiments/graph_trainer/common_utils.py +++ b/torchtitan/experiments/graph_trainer/common_utils.py @@ -419,15 +419,18 @@ def apply_simple_fsdp( (the routed-expert weights) are separately wrapped on the EDP mesh when expert parallelism is enabled. """ + use_local_compute = parallel_dims.spmd_backend == "spmd_types" + compute_mesh = parallel_dims.spmd_dense_mesh() if use_local_compute else None + fsdp_axis = "dp_shard" if use_local_compute else "fsdp" if parallel_dims.dp_replicate_enabled: if parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled: - dp_mesh_dim_names = ["dp_replicate", "fsdp"] + dp_mesh_dim_names = ["dp_replicate", fsdp_axis] dp_mode = "hybrid_shard" else: dp_mesh_dim_names = ["dp_replicate"] dp_mode = "replicate" else: - dp_mesh_dim_names = ["fsdp"] + dp_mesh_dim_names = [fsdp_axis] dp_mode = "fully_shard" dp_mesh = parallel_dims.get_mesh(dp_mesh_dim_names) @@ -462,6 +465,8 @@ def apply_simple_fsdp( dp_mode, mp_policy=mp_policy, shard_dim=experts_shard_dim, + local_compute=use_local_compute, + compute_mesh=compute_mesh, ) model = data_parallel( @@ -469,6 +474,8 @@ def apply_simple_fsdp( dp_mesh, dp_mode, mp_policy=mp_policy, + local_compute=use_local_compute, + compute_mesh=compute_mesh, ) logger.info( "Applied Data Parallel (simple_fsdp) (dp mode=%s) to the model", dp_mode diff --git a/torchtitan/experiments/graph_trainer/dynamic_shapes.py b/torchtitan/experiments/graph_trainer/dynamic_shapes.py index 80403a88ed..1638d53e74 100644 --- a/torchtitan/experiments/graph_trainer/dynamic_shapes.py +++ b/torchtitan/experiments/graph_trainer/dynamic_shapes.py @@ -26,6 +26,8 @@ from typing import Any +import spmd_types as spmd + import torch from torch._subclasses import FakeTensorMode from torch.utils._python_dispatch import is_traceable_wrapper_subclass @@ -183,6 +185,8 @@ def copy_tensor_annotations(fake_arg: torch.Tensor) -> torch.Tensor: for name in _DYNAMO_SHAPE_ANNOTATION_NAMES: if hasattr(arg, name): setattr(fake_arg, name, getattr(arg, name)) + if spmd.has_local_type(arg) or spmd.get_partition_spec(arg) is not None: + spmd.assert_type_like(fake_arg, arg) return fake_arg symbolic_context = _symbolic_context_for_marked_dims(arg) diff --git a/torchtitan/experiments/graph_trainer/llama3/parallelize.py b/torchtitan/experiments/graph_trainer/llama3/parallelize.py index a9bcd83305..8fe83e86a7 100644 --- a/torchtitan/experiments/graph_trainer/llama3/parallelize.py +++ b/torchtitan/experiments/graph_trainer/llama3/parallelize.py @@ -55,7 +55,7 @@ def parallelize_llama( annotate_llama(model) - if parallel_dims.tp_enabled: + if parallel_dims.tp_enabled or parallel_dims.spmd_backend == "spmd_types": model.parallelize(parallel_dims) # Apply simple_fsdp unconditionally. The `fsdp` mesh always exists with a diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py index 835dab86bd..7f3d6536e3 100644 --- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py +++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py @@ -11,9 +11,12 @@ from dataclasses import dataclass from typing import Any +import spmd_types as spmd + import torch import torch.nn as nn import torch.utils._pytree as pytree +from spmd_types.checker import typecheck as spmd_typecheck from torch._guards import tracing, TracingContext from torch._subclasses import FakeTensorMode from torch.fx.experimental.proxy_tensor import make_fx @@ -429,12 +432,14 @@ def _trace_with_args(*args: Any, **kwargs: Any) -> TracedResult: allow_non_fake_inputs=True, shape_env=torch.fx.experimental.symbolic_shapes.ShapeEnv(), ) - fake_args = tuple( - _fakeify_input(fake_mode, a, input_name=f"input_{i}") - if isinstance(a, torch.Tensor) - else a - for i, a in enumerate(unwrapped_args) - ) + typecheck_trace = spmd.is_type_checking() + with spmd.no_typecheck(): + fake_args = tuple( + _fakeify_input(fake_mode, a, input_name=f"input_{i}") + if isinstance(a, torch.Tensor) + else a + for i, a in enumerate(unwrapped_args) + ) output_layouts: dict[int, SubclassLayout] = {} num_flat_outputs: int = 0 @@ -459,9 +464,16 @@ def fn_with_subclass_handling(*plain_args: Any) -> list: if prepared is not None: user_args, user_kwargs = prepared - with _reparametrize_train_state( - module, optimizer, model_state_t, optim_state_t - ), torch.compiler._patch_engine_backward(): + typecheck_context = ( + spmd_typecheck() if typecheck_trace else contextlib.nullcontext() + ) + with ( + _reparametrize_train_state( + module, optimizer, model_state_t, optim_state_t + ), + torch.compiler._patch_engine_backward(), + typecheck_context, + ): result = fn(*user_args, **user_kwargs) flat_outs, output_spec = pytree.tree_flatten(result) @@ -492,6 +504,7 @@ def fn_with_subclass_handling(*plain_args: Any) -> list: _skip_nested_compile(), torch.autograd.set_multithreading_enabled(False), torch.compiler._non_strict_tracing_context(), + spmd.no_typecheck(), ): traced = make_fx( fn_with_subclass_handling, diff --git a/torchtitan/experiments/graph_trainer/simple_fsdp.py b/torchtitan/experiments/graph_trainer/simple_fsdp.py index 19516edd64..a941e411fc 100644 --- a/torchtitan/experiments/graph_trainer/simple_fsdp.py +++ b/torchtitan/experiments/graph_trainer/simple_fsdp.py @@ -9,6 +9,8 @@ from contextlib import contextmanager from dataclasses import dataclass +import spmd_types as spmd + import torch import torch.nn as nn @@ -142,9 +144,12 @@ def _register_parametrization( TODO: In checkpoint saving/loading, avoid parametrization calls when calling get_model_state_dict func in torchtitan's torchtitan/components/checkpoint.py. """ + object.__setattr__(module, "_simple_fsdp_parametrization", parametrization) param_name_to_property = { param_name: property( - lambda self, pn=param_name: parametrization(self._parameters[pn]) + lambda self, pn=param_name: self._simple_fsdp_parametrization( + self._parameters[pn], pn + ) ) for param_name in param_names } @@ -172,6 +177,8 @@ def __init__( mode: str, mp_policy: MixedPrecisionPolicy | None, full_dtensor: bool = False, + local_compute: bool = False, + compute_types: dict[str, tuple[dict, object | None]] | None = None, ) -> None: super().__init__() self.device_mesh = device_mesh @@ -185,6 +192,8 @@ def __init__( self.param_dtype: torch.dtype | None = mp_policy.param_dtype self.reduce_dtype: torch.dtype | None = mp_policy.reduce_dtype self.full_dtensor = full_dtensor + self.local_compute = local_compute + self.compute_types = compute_types def replicate_compute(self, x: DTensor) -> torch.Tensor: # data parallel runtime replicate parameters and do local compute @@ -225,8 +234,12 @@ def replicate_compute(self, x: DTensor) -> torch.Tensor: ) non_dp_mesh = x._spec.mesh[non_dp_mesh_dim_names] - output = DTensor.from_local( - replicated_local_tensor, non_dp_mesh, non_dp_placements + output = ( + replicated_local_tensor + if self.local_compute + else DTensor.from_local( + replicated_local_tensor, non_dp_mesh, non_dp_placements + ) ) elif non_dp_mesh_dims == 0: output = x.redistribute( @@ -244,7 +257,7 @@ def replicate_compute(self, x: DTensor) -> torch.Tensor: return output - def forward(self, x: DTensor) -> torch.Tensor: + def forward(self, x: DTensor, param_name: str) -> torch.Tensor: global _active_parametrization # This should never be set to true during forward, only outside for model # inspection / debugging / initialization @@ -254,7 +267,15 @@ def forward(self, x: DTensor) -> torch.Tensor: if not _active_parametrization: return x - output = self.replicate_compute(x) + if self.local_compute: + with spmd.no_typecheck(): + output = self.replicate_compute(x) + if spmd.is_type_checking(): + assert self.compute_types is not None + local_type, partition_spec = self.compute_types[param_name] + spmd.assert_type(output, local_type, partition_spec=partition_spec) + else: + output = self.replicate_compute(x) return output @@ -265,6 +286,8 @@ def data_parallel( mp_policy: MixedPrecisionPolicy | None = None, shard_dim: int = 0, full_dtensor: bool = False, + local_compute: bool = False, + compute_mesh: DeviceMesh | None = None, ) -> nn.Module: param_sharding: tuple[Placement, ...] if mode == "replicate": @@ -289,16 +312,52 @@ def data_parallel( if "SimpleFSDP" in mod.__class__.__name__: continue + compute_types = ( + { + p_name: (dict(spmd.get_local_type(p)), spmd.get_partition_spec(p)) + for p_name, p in params_dict.items() + if p is not None and p.numel() > 0 + } + if local_compute + else None + ) + for p_name, p in params_dict.items(): if p is not None and p.numel() > 0: + if ( + local_compute + and compute_mesh is not None + and compute_mesh["tp"].size() > 1 + and not isinstance(p, DTensor) + ): + tp_mesh = compute_mesh["tp"] + tp_axis = spmd.normalize_axis(tp_mesh.get_group()) + tp_type = spmd.get_axis_local_type(p, tp_axis) + if tp_type is spmd.V: + partition_spec = spmd.get_partition_spec(p) + assert partition_spec is not None + tp_type = next( + spmd.S(dim) + for dim, entry in enumerate(partition_spec) + if tp_axis + in (entry if isinstance(entry, tuple) else (entry,)) + ) + p = DTensor.from_local( + p, + tp_mesh, + (spmd.spmd_type_to_dtensor_placement(tp_type),), + run_check=False, + ) distribute_tensor_func = ( _distribute_dtensor if isinstance(p, DTensor) else distribute_tensor ) + distributed_param = distribute_tensor_func( + p, device_mesh, param_sharding + ) + registered_param = nn.Parameter(distributed_param) mod.register_parameter( p_name, - nn.Parameter( - distribute_tensor_func(p, device_mesh, param_sharding) - ), + registered_param, ) # to be compatible with DCP, we use a customized _register_parametrization @@ -324,6 +383,8 @@ def data_parallel( mode, mp_policy=mp_policy, full_dtensor=full_dtensor, + local_compute=local_compute, + compute_types=compute_types, ), ) return model diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py index 6fc60f1919..ef2d930522 100644 --- a/torchtitan/experiments/graph_trainer/trainer.py +++ b/torchtitan/experiments/graph_trainer/trainer.py @@ -8,6 +8,8 @@ from dataclasses import dataclass, field from typing import Any +import spmd_types as spmd + import torch import torch.nn as nn @@ -235,7 +237,7 @@ def _make_fx_forward_backward_step( passes, compile_config=self.config.compile, ) - with self.train_context(): + with self.train_context(), spmd.no_typecheck(): outputs = run_traced(self._traced_step, module=model)( inputs, labels,