diff --git a/.github/workflows/test_torchtitan.yml b/.github/workflows/test_torchtitan.yml index 003c7b9b..a80ee4cc 100644 --- a/.github/workflows/test_torchtitan.yml +++ b/.github/workflows/test_torchtitan.yml @@ -13,7 +13,7 @@ concurrency: jobs: test-torchtitan: - name: Test TorchTitan Integration (cuda12.6-py3.12) + name: Test TorchTitan Integration (cuda13.0-py3.12) uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main strategy: fail-fast: true @@ -21,9 +21,9 @@ jobs: include: - name: 12xlargegpu runs-on: linux.g5.12xlarge.nvidia.gpu - torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cu126' + torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cu130' gpu-arch-type: "cuda" - gpu-arch-version: "12.6" + gpu-arch-version: "13.0" with: timeout: 60 runner: ${{ matrix.runs-on }} @@ -38,24 +38,16 @@ jobs: pip install --quiet -r requirements-test.txt # For some reason the spec above isnt working pip uninstall -y torch - pip install --no-input --quiet --pre torch --index-url https://download.pytorch.org/whl/nightly/cu126 + pip install --no-input --quiet --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130 pip install --quiet . # Clone TorchTitan git clone https://github.com/pytorch/torchtitan.git cd torchtitan + git fetch origin pull/3308/head:remove-moe-for-loop-fallback + git checkout remove-moe-for-loop-fallback pip install --quiet -r requirements.txt - # Run TorchTitan training with AutoParallel - NGPU=4 ./run_train.sh \ - --module autoparallel.llama3 \ - --config autoparallel_llama3_debugmodel \ - --parallelism.tensor_parallel_degree 4 - - # TODO: Re-enable deepseek_v3 test once torchtitan experiment is fixed - # (deepseek_v3 experiment is also disabled in torchtitan's own CI) - # NGPU=4 ./run_train.sh \ - # --module autoparallel.deepseek_v3 \ - # --config autoparallel_deepseek_v3_debugmodel \ - # --parallelism.data_parallel_shard_degree 4 \ - # --parallelism.expert_parallel_degree 4 + # Check that AutoParallel and TorchTitan DeepSeek V3 produce matching + # distributed loss and gradient norms for the same 4-GPU debug shape. + torchrun --standalone --nproc-per-node 4 ../tests/torchtitan_dsv3_equivalence.py diff --git a/autoparallel/_testing/models/dsv3.py b/autoparallel/_testing/models/dsv3.py index 5a897b71..99b41cba 100644 --- a/autoparallel/_testing/models/dsv3.py +++ b/autoparallel/_testing/models/dsv3.py @@ -14,10 +14,10 @@ import triton.language as tl from torch import nn from torch.distributed.tensor import DeviceMesh, DTensor -from torch.distributed.tensor.placement_types import Partial, Replicate, Shard +from torch.distributed.tensor.placement_types import Replicate, Shard from torch.nn.attention import SDPBackend, sdpa_kernel -from autoparallel.collectives import all_to_all, axis_size, local_map +from autoparallel.collectives import all_reduce, all_to_all, axis_size, local_map _MODULE_FQN = "module_fqn" @@ -428,12 +428,12 @@ def forward( # top_scores is still derived from the original scores. if expert_bias is not None: _, selected_experts_indices = torch.topk( - scores + expert_bias, k=self.top_k, dim=-1 + scores + expert_bias, k=self.top_k, dim=-1, sorted=False ) top_scores = scores.gather(dim=-1, index=selected_experts_indices) else: top_scores, selected_experts_indices = torch.topk( - scores, k=self.top_k, dim=-1 + scores, k=self.top_k, dim=-1, sorted=False ) if self.score_func == "sigmoid" and self.route_norm: @@ -616,7 +616,6 @@ def _token_combine( return routed_output -# @torch.library.custom_op("autoparallel::local_mapped_region", mutates_args=()) def local_mapped_region( x: torch.Tensor, selected_experts_indices: torch.Tensor, @@ -629,12 +628,10 @@ def local_mapped_region( num_experts: int, score_before_experts: bool, axis_name: str, + token_count_reduce_axis_names: tuple[str, ...], ) -> tuple[torch.Tensor, torch.Tensor]: - # assert False, f"{x.shape}, {selected_experts_indices.shape}, {top_scores.shape}, {out.shape}" - dim = x.shape[-1] - # num_tokens_per_expert = torch.ops.autoparallel.batched_histc( num_tokens_per_expert = torch.histc( selected_experts_indices.flatten(), bins=num_experts, @@ -642,8 +639,15 @@ def local_mapped_region( max=num_experts, ) - # total_tokens_per_expert = all_reduce(num_tokens_per_expert, axis_name) + # Token counts are logically Partial(sum), Partial(sum). TorchTitan's + # optimizer hook aggregates them with an all-reduce, so reduce here and mark + # the local_map output replicated to match the optimizer-side contract. total_tokens_per_expert = num_tokens_per_expert + for axis_name_for_token_counts in token_count_reduce_axis_names: + total_tokens_per_expert = all_reduce( + total_tokens_per_expert, + axis_name_for_token_counts, + ) token_indices_experts_sorted = torch.argsort( selected_experts_indices.view(-1), stable=True @@ -703,98 +707,6 @@ def local_mapped_region( return out, total_tokens_per_expert -# @local_mapped_region.register_fake -def _( - routed_input: torch.Tensor, - selected_expert_indices: torch.Tensor, - top_scores: torch.Tensor, - out: torch.Tensor, - experts_w1: torch.Tensor, - experts_w2: torch.Tensor, - experts_w3: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - num_experts = 64 - return torch.empty_like(routed_input), torch.empty( - (1, num_experts), dtype=routed_input.dtype, device=routed_input.device - ) - - -# @torch.library.custom_op("autoparallel::local_mapped_region_grad", mutates_args=()) -# def local_mapped_region_grad( -# routed_input: torch.Tensor, -# selected_experts_indices: torch.Tensor, -# top_scores: torch.Tensor, -# out: torch.Tensor, -# experts_w1: torch.Tensor, -# experts_w2: torch.Tensor, -# experts_w3: torch.Tensor, -# ) -> tuple[ -# torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor -# ]: -# grad_i = torch.empty_like(routed_input) -# grad_o = torch.empty_like(out) -# grad_s = torch.empty_like(top_scores) -# g1 = torch.empty_like(experts_w1) -# g2 = torch.empty_like(experts_w2) -# g3 = torch.empty_like(experts_w3) -# return grad_i, grad_s, grad_o, g1, g2, g3 - - -# @local_mapped_region_grad.register_fake -# def _( -# routed_input: torch.Tensor, -# selected_experts_indices: torch.Tensor, -# top_scores: torch.Tensor, -# out: torch.Tensor, -# experts_w1: torch.Tensor, -# experts_w2: torch.Tensor, -# experts_w3: torch.Tensor, -# ) -> tuple[ -# torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor -# ]: -# grad_i = torch.empty_like(routed_input) -# grad_o = torch.empty_like(out) -# grad_s = torch.empty_like(top_scores) -# g1 = torch.empty_like(experts_w1) -# g2 = torch.empty_like(experts_w2) -# g3 = torch.empty_like(experts_w3) -# return grad_i, grad_s, grad_o, g1, g2, g3 - - -# def setup_context_local_mapped_region(ctx, inputs, output): -# # routed_input, num_tokens_per_expert, experts_w1, experts_w2, experts_w3 = inputs -# ctx.save_for_backward(*inputs) - - -# def backward_local_mapped_region(ctx, grad, grad2): -# ( -# routed_input, -# selected_experts_indices, -# top_scores, -# out, -# experts_w1, -# experts_w2, -# experts_w3, -# ) = ctx.saved_tensors -# grad_i, grad_s, grad_o, g1, g2, g3 = local_mapped_region_grad( -# routed_input, -# selected_experts_indices, -# top_scores, -# out, -# experts_w1, -# experts_w2, -# experts_w3, -# ) -# return grad_i, None, grad_s, grad_o, g1, g2, g3 - - -# torch.library.register_autograd( -# "autoparallel::local_mapped_region", -# backward_local_mapped_region, -# setup_context=setup_context_local_mapped_region, -# ) - - def _moe_forward( x: torch.Tensor, router_gate_weight: torch.Tensor, @@ -812,56 +724,16 @@ def _moe_forward( score_before_experts: bool, compute_dtype: torch.dtype | None = None, ): - # x: 64, 2048, 256 bs, slen, dim = x.shape x = x.view(-1, dim) - # top_scores and selected_experts_indices shape (bs*slen*top_k,) - # num_tokens_per_expert shape (num_experts,) ( top_scores, selected_experts_indices, ) = router(x, router_gate_weight, expert_bias) - # tokens_per_expert will be used to update the expert bias for load balancing. - # and also to count the expert usage - # TODO: Activation Checkpointing has the side effect of double counting tokens_per_expert -- - # first in the forward pass, and then in the backward pass. However, this has no - # effect on the expert bias update thanks to the torch.sign() operator. - # moved out to remove mutation - # with torch.no_grad(): - # tokens_per_expert.add_(num_tokens_per_expert) - - # top_scores and token_indices_experts_sorted shape (bs*slen*top_k,) - # num_tokens_per_expert shape (num_experts,) - # NOTE: the reason we need to compute num_tokens_per_expert again is: - # 1st computation in router is to update self.tokens_per_expert - # which would be the same across all TP ranks. - # 2nd computation in reorderer is for the actual routing and experts computation - # which would be sharded over TP ranks if expert_tensor_parallel_degree==1. - # If tensor_paralllel_degree == expert_tensor_parallel_degree, they agree. - # ( - # top_scores_experts_sorted, - # token_indices_experts_sorted, - # # _, #num_tokens_per_expert, - # ) = reorderer(top_scores, selected_experts_indices) - - # shape (bs*slen*top_k, dim) - # routed_output = experts(routed_input, num_tokens_per_expert) - out = functional_feed_forward(shared_w1, shared_w2, shared_w3, x, compute_dtype) - ###################################################### - # This is in the local_map region - ###################################################### - - # expert_placements = ((Replicate(), Shard(0)),) * 3 - # in_placements = ( - # (Shard(0), Shard(0)), - # (Shard(0), Shard(0)), - # (Shard(0), Shard(0)), - # (Shard(0), Shard(0)), - # ) # Dynamo reorders captured variables (lifted freevars) before explicit # arguments, so x must come first in the input order and placements. reordered_placements = ( @@ -876,13 +748,18 @@ def _moe_forward( None, None, None, + None, ) + token_count_reduce_axis_names: tuple[str, ...] = () + if mesh is not None and mesh.mesh_dim_names is not None: + token_count_reduce_axis_names = tuple(mesh.mesh_dim_names) + out, num_tokens_per_expert = local_map( local_mapped_region, out_placements=( (Shard(0), Shard(0)), - (Partial(reduce_op="sum"), Partial(reduce_op="sum")), + (Replicate(), Replicate()), ), in_placements=reordered_placements, redistribute_inputs=True, @@ -900,20 +777,9 @@ def _moe_forward( router.num_experts, score_before_experts, axis_name, + token_count_reduce_axis_names, ) - # assert False, f"there: {out.shape}, {num_tokens_per_expert.shape}" - - ###################################################### - # end of the local_map region - ###################################################### - - # shared expert - # if shared_experts is not None: - # out = shared_experts(x) - # else: - # out = torch.zeros_like(x) - # assert False, f"{out.shape}, {token_indices_experts_sorted.shape}, {routed_output.shape}" out = out.reshape(bs, slen, dim) return out, num_tokens_per_expert @@ -1003,7 +869,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self.compute_dtype, ) - # HOPs don't support buffer mutations, keep this outside + # HOPs don't support buffer mutations, keep this outside. with torch.no_grad(): self.tokens_per_expert.add_(num_tokens_per_expert) # type: ignore[operator] return out @@ -1049,7 +915,6 @@ def _init_backend(cls) -> None: if cls.backends: return - # Add CuDNN on B200 w/ highest priority cls.backends = [ SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, @@ -1526,22 +1391,16 @@ def forward( ) # (bsz, seqlen, dim) def init_weights(self, init_std: float): - linear_list = [ - self.wkv_a, - self.wkv_b, - ] if self.q_lora_rank > 0: - linear_list.extend([self.wq_a, self.wq_b]) + nn.init.trunc_normal_(self.wq_a.weight, mean=0.0, std=0.02) + self.q_norm.reset_parameters() + nn.init.trunc_normal_(self.wq_b.weight, mean=0.0, std=0.02) else: - linear_list.append(self.wq) - - for linear in linear_list: - nn.init.trunc_normal_(linear.weight, mean=0.0, std=0.02) - nn.init.trunc_normal_(self.wo.weight, mean=0.0, std=init_std) - + nn.init.trunc_normal_(self.wq.weight, mean=0.0, std=0.02) + nn.init.trunc_normal_(self.wkv_a.weight, mean=0.0, std=0.02) self.kv_norm.reset_parameters() - if self.q_lora_rank > 0: - self.q_norm.reset_parameters() + nn.init.trunc_normal_(self.wkv_b.weight, mean=0.0, std=0.02) + nn.init.trunc_normal_(self.wo.weight, mean=0.0, std=init_std) class TransformerBlock(nn.Module): @@ -1581,7 +1440,7 @@ def __init__( route_norm=moe_cfg.router.route_norm, route_scale=moe_cfg.router.route_scale, score_before_experts=moe_cfg.experts.token_dispatcher.score_before_experts, - use_grouped_mm=moe_cfg.experts.use_grouped_mm, + use_grouped_mm=getattr(moe_cfg.experts, "use_grouped_mm", True), load_balance_coeff=moe_cfg.load_balance_coeff, mesh=mesh, compute_dtype=compute_dtype, @@ -1621,9 +1480,9 @@ def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor): return x def init_weights(self, buffer_device: torch.device): - for norm in (self.attention_norm, self.ffn_norm): - norm.reset_parameters() self.attention.init_weights(self.weight_init_std) + self.attention_norm.reset_parameters() + self.ffn_norm.reset_parameters() if self.moe_enabled: self.moe.init_weights(self.weight_init_std, buffer_device) else: @@ -1673,8 +1532,10 @@ def __init__( def init_weights( self, buffer_device: torch.device | None = None, seed: int | None = None ) -> None: - _init_weights_tok_embeddings(self, seed) - _init_weights_layers(self, buffer_device, seed) + if seed is not None: + torch.manual_seed(seed) + _init_weights_tok_embeddings(self) + _init_weights_layers(self, buffer_device) _init_weights_norm_and_output(self) def forward( @@ -1722,9 +1583,7 @@ def forward( return output -def _init_weights_tok_embeddings(self: DeepSeekV3Model, seed: int | None = None): - if seed is not None: - torch.manual_seed(seed) +def _init_weights_tok_embeddings(self: DeepSeekV3Model): if self.tok_embeddings is not None: nn.init.normal_(self.tok_embeddings.weight) @@ -1732,15 +1591,12 @@ def _init_weights_tok_embeddings(self: DeepSeekV3Model, seed: int | None = None) def _init_weights_layers( self: DeepSeekV3Model, buffer_device: torch.device | None, - seed: int | None = None, ): if buffer_device is None: buffer_device = self.freqs_cis.device # type: ignore[assignment] with torch.device(buffer_device): # type: ignore[arg-type] self.freqs_cis = precompute_freqs_cis(self.model_args) for i, layer in enumerate(self.layers.values()): - if seed is not None: - torch.manual_seed(seed) if layer is not None: assert isinstance(layer, TransformerBlock) layer.init_weights(buffer_device) # type: ignore[arg-type] diff --git a/autoparallel/api.py b/autoparallel/api.py index 932dc5f2..dea814ab 100644 --- a/autoparallel/api.py +++ b/autoparallel/api.py @@ -389,7 +389,12 @@ def optimize_placement(self, verbose=True): return self.sharding_placement - def _apply_placement_common(self, sharding_placement): + def _apply_placement_common( + self, + sharding_placement, + *, + decompose_after_sharding=True, + ): t0 = time.perf_counter() self._assert_entered() @@ -422,6 +427,7 @@ def _apply_placement_common(self, sharding_placement): sharding_placement, self.joint_with_descriptors.params_spec, self.joint_with_descriptors.buffers_spec, + decompose_after_sharding=decompose_after_sharding, ) t_apply = time.perf_counter() # clean it up by removing the added aliases from previous pass @@ -475,9 +481,15 @@ def _apply_placement_common(self, sharding_placement): sharded_buffer_dict, ) - def apply_placement(self, sharding_placement): + def apply_placement( + self, + sharding_placement, + *, + decompose_after_sharding=True, + ): sharded_param_dict, sharded_buffer_dict = self._apply_placement_common( - sharding_placement + sharding_placement, + decompose_after_sharding=decompose_after_sharding, ) mark_fsdp_all_gather_recomputation( diff --git a/autoparallel/apply_sharding.py b/autoparallel/apply_sharding.py index 21ba3b40..9c4a78e8 100644 --- a/autoparallel/apply_sharding.py +++ b/autoparallel/apply_sharding.py @@ -419,10 +419,15 @@ def _make_local_args(gm, sharding_placement): return local_args -def _lower_to_parallel_graph(gm, sharding_placement, local_args, dynamic=False): +def _lower_to_parallel_graph( + gm, + sharding_placement, + local_args, + dynamic=False, + *, + decompose_after_sharding=True, +): """Two-pass lowering: interpret with sharding collectives, then decompose.""" - decomp_table = _get_inductor_decomp_table() - interp = ApplyShardingInterpreter(gm, sharding_placement, dynamic=dynamic) tracing_mode = "symbolic" if dynamic else "real" @@ -432,12 +437,11 @@ def _lower_to_parallel_graph(gm, sharding_placement, local_args, dynamic=False): cleanup_graph(parallel_gm0) interp2 = torch.fx.Interpreter(parallel_gm0) + make_fx_kwargs = {"tracing_mode": tracing_mode} + if decompose_after_sharding: + make_fx_kwargs["decomposition_table"] = _get_inductor_decomp_table() with fx_traceback.preserve_node_meta(): - parallel_gm = make_fx( - interp2.run, - decomposition_table=decomp_table, - tracing_mode=tracing_mode, - )(*local_args) + parallel_gm = make_fx(interp2.run, **make_fx_kwargs)(*local_args) cleanup_graph(parallel_gm) return parallel_gm @@ -483,7 +487,14 @@ def _shard_params_and_buffers(gm, sharding_placement, params_spec, buffers_spec) return sharded_param_dict, sharded_buffer_dict -def apply_sharding_to_model(gm, sharding_placement, params_spec, buffers_spec): +def apply_sharding_to_model( + gm, + sharding_placement, + params_spec, + buffers_spec, + *, + decompose_after_sharding=True, +): t0 = time.perf_counter() dynamic = _has_symbolic_shapes(gm) @@ -516,7 +527,13 @@ def apply_sharding_to_model(gm, sharding_placement, params_spec, buffers_spec): local_args = _make_local_args(gm, sharding_placement) t1 = time.perf_counter() - parallel_gm = _lower_to_parallel_graph(gm, sharding_placement, local_args, dynamic) + parallel_gm = _lower_to_parallel_graph( + gm, + sharding_placement, + local_args, + dynamic, + decompose_after_sharding=decompose_after_sharding, + ) t2 = time.perf_counter() _copy_descriptors_and_rename_placeholders(gm, parallel_gm) diff --git a/tests/torchtitan_dsv3_equivalence.py b/tests/torchtitan_dsv3_equivalence.py new file mode 100644 index 00000000..4796de5f --- /dev/null +++ b/tests/torchtitan_dsv3_equivalence.py @@ -0,0 +1,445 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. +# +# This source code is licensed under the BSD license found in the +# LICENSE file in the root directory of this source tree. + +"""Compare AutoParallel's local_map DSv3 model with TorchTitan's DSv3 model. + +This is a 4-GPU distributed integration check. It verifies that both models can +load the same full state exactly, then compares one forward/backward step. +""" + +import os +import sys +from pathlib import Path + +import torch +import torch.distributed as dist +from torch.distributed.fsdp import MixedPrecisionPolicy +from torch.distributed.tensor import DTensor +from torch.distributed.tensor.placement_types import Replicate, Shard + +from autoparallel._testing.models.dsv3 import ( + DeepSeekV3Model as AutoParallelDeepSeekV3Model, +) +from autoparallel._testing.models.dsv3 import make_dsv3_config +from autoparallel.api import AutoParallel + +WORLD_SIZE = 4 +LOCAL_BATCH_SIZE = 2 +SEQ_LEN = 1024 +SEED = 123 + +DSV3_DIM = 256 +DSV3_VOCAB_SIZE = 2048 +DSV3_ROPE_DIM = 64 +DSV3_NUM_LAYERS = 4 +DSV3_NUM_DENSE_LAYERS = 0 +DSV3_NUM_EXPERTS = 4 +DSV3_NUM_SHARED_EXPERTS = 2 +DSV3_ROUTER_TOP_K = 2 +DSV3_DENSE_HIDDEN_DIM = 1024 +DSV3_MOE_HIDDEN_DIM = 256 + +AP_DP_DEGREE = 2 +AP_EP_DEGREE = 2 +TT_DP_SHARD_DEGREE = 4 +TT_EP_DEGREE = 2 + +NUMERICS_RTOL = 5e-4 +NUMERICS_ATOL = 1e-4 + + +def _import_torchtitan(): + try: + import torchtitan # noqa: F401 + except ModuleNotFoundError: + candidates = ( + # AutoParallel CI runs this script from inside a cloned TorchTitan + # checkout. + Path.cwd(), + # Local development often keeps TorchTitan next to AutoParallel. + Path(__file__).resolve().parents[2] / "torchtitan", + ) + for candidate in candidates: + if (candidate / "torchtitan").exists(): + sys.path.insert(0, str(candidate)) + break + import torchtitan # noqa: F401 + + +def _make_autoparallel_config(seq_len: int): + return make_dsv3_config( + num_experts=DSV3_NUM_EXPERTS, + top_k=DSV3_ROUTER_TOP_K, + n_layers=DSV3_NUM_LAYERS, + n_dense_layers=DSV3_NUM_DENSE_LAYERS, + max_seq_len=seq_len, + ) + + +def _build_torchtitan_config(seq_len: int): + _import_torchtitan() + + from torchtitan.models.deepseek_v3 import ( + _EMBEDDING_INIT, + _NORM_INIT, + DeepSeekV3Model, + Embedding, + Linear, + RMSNorm, + RoPE, + _build_dsv3_layers, + _output_linear_init, + ) + + layers = _build_dsv3_layers( + n_layers=DSV3_NUM_LAYERS, + n_dense_layers=DSV3_NUM_DENSE_LAYERS, + dim=DSV3_DIM, + n_heads=16, + q_lora_rank=0, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=DSV3_ROPE_DIM, + v_head_dim=128, + mscale=0.70, + dense_hidden_dim=DSV3_DENSE_HIDDEN_DIM, + moe_hidden_dim=DSV3_MOE_HIDDEN_DIM, + num_experts=DSV3_NUM_EXPERTS, + num_shared_experts=DSV3_NUM_SHARED_EXPERTS, + router_top_k=DSV3_ROUTER_TOP_K, + router_score_func="softmax", + score_before_experts=False, + attn_backend="sdpa", + moe_comm_backend="standard", + non_blocking_capacity_factor=None, + ) + for layer_config in layers: + layer_config.attention.rope_max_seq_len = seq_len + + return DeepSeekV3Model.Config( + vocab_size=DSV3_VOCAB_SIZE, + dim=DSV3_DIM, + tok_embeddings=Embedding.Config( + num_embeddings=DSV3_VOCAB_SIZE, + embedding_dim=DSV3_DIM, + param_init=_EMBEDDING_INIT, + ), + norm=RMSNorm.Config(normalized_shape=DSV3_DIM, param_init=_NORM_INIT), + lm_head=Linear.Config( + in_features=DSV3_DIM, + out_features=DSV3_VOCAB_SIZE, + param_init=_output_linear_init(DSV3_DIM), + ), + rope=RoPE.Config( + dim=DSV3_ROPE_DIM, + max_seq_len=seq_len, + theta=10000.0, + backend="complex", + scaling="yarn", + rope_factor=40.0, + beta_fast=32.0, + beta_slow=1.0, + original_seq_len=4096, + ), + layers=layers, + ) + + +def _init_distributed() -> torch.device: + if "WORLD_SIZE" not in os.environ: + raise RuntimeError(f"run with torchrun --nproc-per-node {WORLD_SIZE}") + if int(os.environ["WORLD_SIZE"]) != WORLD_SIZE: + raise RuntimeError(f"expected exactly {WORLD_SIZE} ranks") + local_rank = int(os.environ["LOCAL_RANK"]) + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + dist.init_process_group("nccl", device_id=device) + return device + + +def _make_seed_state(seq_len: int): + config = _make_autoparallel_config(seq_len) + model = AutoParallelDeepSeekV3Model(config, compute_dtype=torch.bfloat16) + model.init_weights(buffer_device=torch.device("cpu"), seed=SEED) + return { + name: tensor.detach().clone() for name, tensor in model.state_dict().items() + } + + +def _copy_full_tensor(target: torch.Tensor, source: torch.Tensor, device: torch.device): + source = source.to(device=device, dtype=target.dtype) + if isinstance(target, DTensor): + replicated_source = DTensor.from_local( + source, + device_mesh=target.device_mesh, + placements=(Replicate(),) * target.device_mesh.ndim, + ) + source = replicated_source.redistribute( + target.device_mesh, + target.placements, + ) + target.copy_(source) + + +def _load_full_state( + model: torch.nn.Module, + state: dict[str, torch.Tensor], + device: torch.device, +): + targets: dict[str, torch.Tensor] = {} + targets.update(model.named_parameters()) + targets.update(model.named_buffers()) + with torch.no_grad(): + for name, source in state.items(): + assert name in targets, f"missing state target {name}" + _copy_full_tensor(targets[name], source, device) + + +def _loss_and_grad_norm(model, tokens: torch.Tensor, labels: torch.Tensor): + logits = model(tokens) + if isinstance(logits, DTensor): + logits = logits.to_local() + assert not torch.any(torch.isnan(logits)), "forward produced NaNs" + + local_loss_sum = torch.nn.functional.cross_entropy( + logits.flatten(0, 1).float(), + labels.flatten(0, 1), + reduction="sum", + ) + global_loss_sum = local_loss_sum.detach().clone() + dist.all_reduce(global_loss_sum, op=dist.ReduceOp.SUM) + global_tokens = torch.tensor( + labels.numel() * dist.get_world_size(), + device=labels.device, + dtype=torch.float32, + ) + loss = local_loss_sum / global_tokens + loss.backward() + + local_grad_sq = torch.zeros((), device=labels.device) + for parameter in model.parameters(): + grad = parameter.grad + if grad is None: + continue + if isinstance(grad, DTensor): + grad = grad.to_local() + local_grad_sq = local_grad_sq + grad.float().pow(2).sum() + dist.all_reduce(local_grad_sq, op=dist.ReduceOp.SUM) + return global_loss_sum / global_tokens, local_grad_sq.sqrt() + + +def _state_for_compare(model): + state = {} + for name, tensor in model.state_dict().items(): + if isinstance(tensor, DTensor): + tensor = tensor.full_tensor() + state[name] = tensor.detach().cpu() + return state + + +def _make_inputs(device: torch.device, local_batch_size: int, seq_len: int): + rank = dist.get_rank() + world_size = dist.get_world_size() + if rank == 0: + full_tokens = torch.randint( + 0, + DSV3_VOCAB_SIZE, + (local_batch_size * world_size, seq_len), + device=device, + ) + else: + full_tokens = torch.empty( + (local_batch_size * world_size, seq_len), + dtype=torch.long, + device=device, + ) + dist.broadcast(full_tokens, src=0) + return full_tokens.chunk(world_size, dim=0)[rank].contiguous() + + +def _run_autoparallel( + device: torch.device, + tokens: torch.Tensor, + seed_state: dict[str, torch.Tensor], +): + seq_len = tokens.shape[1] + mesh = torch.distributed.device_mesh.init_device_mesh( + "cuda", + (AP_DP_DEGREE, AP_EP_DEGREE), + mesh_dim_names=("dp", "ep"), + ) + config = _make_autoparallel_config(seq_len) + global_batch_size = tokens.shape[0] * dist.get_world_size() + + with torch.device("meta"): + model = AutoParallelDeepSeekV3Model( + config, + mesh=mesh, + compute_dtype=torch.bfloat16, + ) + + def input_fn(): + return torch.randint( + 0, + config.vocab_size, + (global_batch_size, seq_len), + device=device, + ) + + mp_policy = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + ) + with AutoParallel( + model, input_fn, mesh, mp_policy=mp_policy, dynamic=True + ) as autop: + autop.add_parameter_memory_constraint(low=None, high=None) + x_sharding = (Shard(0), Shard(0)) + autop.add_input_constraints([x_sharding]) + autop.add_output_constraints([x_sharding]) + parallel_model = autop.apply_placement( + autop.optimize_placement(verbose=False), + decompose_after_sharding=False, + ) + + parallel_model.to_empty(device=device) + _load_full_state(parallel_model, seed_state, device) + loss, grad_norm = _loss_and_grad_norm(parallel_model, tokens, tokens) + return loss, grad_norm, _state_for_compare(parallel_model) + + +def _run_torchtitan( + device: torch.device, + tokens: torch.Tensor, + seed_state: dict[str, torch.Tensor], +): + _import_torchtitan() + + from torchtitan.config import ( + ActivationCheckpointConfig, + CompileConfig, + ParallelismConfig, + TrainingConfig, + ) + from torchtitan.distributed import ParallelDims + from torchtitan.models.deepseek_v3 import DeepSeekV3Model + from torchtitan.models.deepseek_v3.parallelize import parallelize_deepseekv3 + + seq_len = tokens.shape[1] + parallel_dims = ParallelDims( + dp_replicate=1, + dp_shard=TT_DP_SHARD_DEGREE, + cp=1, + tp=1, + pp=1, + ep=TT_EP_DEGREE, + world_size=WORLD_SIZE, + ) + training = TrainingConfig( + local_batch_size=tokens.shape[0], + seq_len=seq_len, + mixed_precision_param="bfloat16", + mixed_precision_reduce="float32", + ) + parallelism = ParallelismConfig( + data_parallel_shard_degree=TT_DP_SHARD_DEGREE, + expert_parallel_degree=TT_EP_DEGREE, + disable_loss_parallel=True, + ) + compile_config = CompileConfig(enable=False) + ac_config = ActivationCheckpointConfig(mode="none") + config = _build_torchtitan_config(seq_len) + config.update_from_config( + trainer_config=type( + "TrainerConfig", + (), + { + "training": training, + "parallelism": parallelism, + "debug": type("DebugConfig", (), {"moe_force_load_balance": False})(), + }, + )() + ) + + with torch.device("meta"): + model = DeepSeekV3Model(config) + parallelize_deepseekv3( + model, + parallel_dims=parallel_dims, + training=training, + parallelism=parallelism, + compile_config=compile_config, + ac_config=ac_config, + dump_folder="/tmp", + ) + model.to_empty(device=device) + _load_full_state(model, seed_state, device) + loss, grad_norm = _loss_and_grad_norm(model, tokens, tokens) + return loss, grad_norm, _state_for_compare(model) + + +def main(): + device = _init_distributed() + try: + torch.manual_seed(SEED) + tokens = _make_inputs( + device, + local_batch_size=LOCAL_BATCH_SIZE, + seq_len=SEQ_LEN, + ) + seed_state = _make_seed_state(seq_len=tokens.shape[1]) + + ap_loss, ap_grad_norm, ap_state = _run_autoparallel( + device, + tokens, + seed_state, + ) + tt_loss, tt_grad_norm, tt_state = _run_torchtitan( + device, + tokens, + seed_state, + ) + + if dist.get_rank() == 0: + for name in sorted(set(ap_state) | set(tt_state)): + if name not in ap_state or name not in tt_state: + raise AssertionError( + f"state missing {name}: " + f"ap={name in ap_state}, tt={name in tt_state}" + ) + if not torch.equal(ap_state[name], tt_state[name]): + diff = (ap_state[name].float() - tt_state[name].float()).abs().max() + raise AssertionError( + f"state differs for {name} with max diff {diff.item()}" + ) + + # The full states above must match exactly. The optimized AP graph may + # still choose sharded matmuls/reductions that differ from TorchTitan's + # FSDP-only arithmetic order, so loss and grad norm are close checks. + torch.testing.assert_close( + ap_loss, + tt_loss, + rtol=NUMERICS_RTOL, + atol=NUMERICS_ATOL, + ) + torch.testing.assert_close( + ap_grad_norm, + tt_grad_norm, + rtol=NUMERICS_RTOL, + atol=NUMERICS_ATOL, + ) + if dist.get_rank() == 0: + print( + "AutoParallel and TorchTitan DSv3 numerics are close: " + f"loss={ap_loss.item():.6f}, grad_norm={ap_grad_norm.item():.6f}" + ) + finally: + dist.barrier(device_ids=[device.index]) + torch.cuda.synchronize(device) + dist.destroy_process_group() + + +if __name__ == "__main__": + main()