garrett361

"""
Minimal reproducer showing that delaying the replicate-dimension all-reduce to
only the final microbatch improves throughput under HSDP with inhomogeneous
microbatch compute costs. Uses variable batch sizes as a stand-in for more realistic
cases, e.g. varying attention costs in packed sequences.

Usage:

    # HSDP
    torchrun --nproc-per-node=4 scripts/bench_delay_all_reduce.py --accum-steps 4

garrett361

Summarizing what I've found about standard model initialization practices in public code bases and papers.

Code Bases

torchtitan (commit 8599818f)

Shared depth-scaling utility (param_init.py): depth_scaled_std(base_std, layer_id) = base_std / sqrt(2 * (layer_id + 1)) (0-indexed layers).

garrett361

❯ torchrun --nproc-per-node 2 test.py
W0225 22:03:36.129000 3702213 torch/distributed/run.py:852]
W0225 22:03:36.129000 3702213 torch/distributed/run.py:852] *****************************************
W0225 22:03:36.129000 3702213 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
W0225 22:03:36.129000 3702213 torch/distributed/run.py:852] *****************************************
[backward] lin1.weight: p.dtype=torch.bfloat16 grad_dtype=torch.float32
[backward] lin0.weight: p.dtype=torch.bfloat16 grad_dtype=torch.bfloat16
[Rank 0] FSDP reduce-scatter expects uniform gradient dtype but got {torch.bfloat16, torch.float32}
[...]
[Rank 1] FSDP reduce-scatter expects uniform gradient dtype but got {torch.float32, torch.bfloat16}

garrett361

Writing out the math used in Sonic MoE: https://arxiv.org/abs/2512.14080. Sec. 3.2 and Appendix C.

Setup

Goal: compute the backwards pass without needing to cache large tensors which would blow up with increased sparsification.

Tensors:

• X_{ted}: input tensors
• W^1_{edn}: up projection

garrett361

R=tensor([[ 0.0673+0.j, -0.6867+0.j,  0.7238+0.j],
        [-0.4553+0.j, -0.6666+0.j, -0.5901+0.j],
        [ 0.8878+0.j, -0.2899+0.j, -0.3575+0.j]])

R@R.T=tensor([[ 1.0000e+00+0.j, -2.7138e-08+0.j, -2.7308e-08+0.j],
        [-2.7138e-08+0.j,  1.0000e+00+0.j, -3.7264e-08+0.j],
        [-2.9802e-08+0.j, -4.4703e-08+0.j,  1.0000e+00+0.j]])

eigvals=tensor([ 1.0000+1.4901e-08j, -0.9785+2.0648e-01j, -0.9785-2.0648e-01j])

garrett361

torchrun --nproc-per-node 2 set_div_err.py
W0605 14:34:10.112000 783116 torch/distributed/run.py:766]
W0605 14:34:10.112000 783116 torch/distributed/run.py:766] *****************************************
W0605 14:34:10.112000 783116 torch/distributed/run.py:766] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
W0605 14:34:10.112000 783116 torch/distributed/run.py:766] *****************************************
Running with trivial mp_policy
Passed on RANK=0 with mp_policy=MixedPrecisionPolicy(param_dtype=None, reduce_dtype=None, output_dtype=None, cast_forward_inputs=True)Passed on RANK=1 with mp_policy=MixedPrecisionPolicy(param_dtype=None, reduce_dtype=None, output_dtype=None, cast_forward_inputs=True)

Running with non-trivial mp_policy
FAILED on RANK=0 with mp_policy=MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16, output

garrett361

torch.__version__='2.8.0.dev20250421+cu126'
Passed first test
Traceback (most recent call last):
  File "/app/torchtitan/gemm_test.py", line 38, in <module>
    test_gemm_backward_fails()
  File "/app/torchtitan/gemm_test.py", line 31, in test_gemm_backward_fails
    out.sum().backward()
  File "/opt/conda/envs/ai/lib/python3.11/site-packages/torch/_tensor.py", line 648, in backward
    torch.autograd.backward(
  File "/opt/conda/envs/ai/lib/python3.11/site-packages/torch/autograd/__init__.py", line 354, in backward

garrett361

import os

import torch
import torch.distributed as dist
from torch.distributed.tensor import Shard, distribute_tensor

if __name__ == "__main__":
    try:
        world_size = int(os.environ["WORLD_SIZE"])
        local_rank = int(os.environ["LOCAL_RANK"])

garrett361

import multiprocessing as mp
import os

import torch
import torch.distributed as dist
from torch.distributed.tensor import distribute_tensor
from torch.distributed.tensor.placement_types import (
    Shard,
)

garrett361

import argparse
import multiprocessing as mp
import os

import torch
import torch.distributed as dist


def print_rank(s: str) -> None:
    s = f'[rank={os.environ["RANK"]}] ' + s