Understanding RMSNorm Gradient Synchronization in Tensor-Parallel LLaMA

[alpemreacar]

I am testing tensor parallelism (TP) using a simple FSDP-based implementation of LLaMA-3. I am on commit 9f211ec, using Python 3.10 and a nightly version of PyTorch, as noted in the README.

My understanding is that under TP, model weights are sharded across ranks so that each rank performs a fraction of the total computation. I am able to capture the execution graphs using TorchInductor by running:

TORCH_COMPILE_DEBUG=1 CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 -m torchtitan.train --job.config_file=tp.oml --model.name simple_fsdp.llama3 --job.custom_config_module=torchtitan.experiments.simple_fsdp.job_config

I have attached my toml configuration file along with a reference implementation.

As expected, I observe both all_gather and reduce_scatter operations in the forward and backward passes, since each rank is computing only a partial result.

However, I am confused about how the RMSNorm parameters are updated. In the LLaMA implementation, each transformer block processes tensors of shape (B, SeqLen / TP, HiddenDim). As a result, each rank only observes gradients corresponding to B × SeqLen / TP tokens for the RMSNorm parameters. These gradients should therefore be reduced across ranks.

However, when inspecting the generated graphs, I do not see any explicit all_reduce corresponding to the RMSNorm parameter gradients.

Could someone clarify where and how this gradient reduction is performed during execution? Is it fused, implicit, or handled outside of the graph I am inspecting?

Any insight would be appreciated.

tptoml.txt

[sanketpurandare]

Understanding RMSNorm Gradient Reduction in Tensor Parallelism

TL;DR: RMSNorm parameter gradients ARE reduced, but through DTensor's Partial placement semantics rather than an explicit all_reduce operation in the compiled graph.

How It Works

1. Sequence Parallel Setup for RMSNorm

RMSNorm is parallelized using SequenceParallel in the layer plan:

# torchtitan/models/llama3/infra/parallelize.py
layer_plan = {
    "attention_norm": SequenceParallel(),
    "ffn_norm": SequenceParallel(),
    # ...
}

In SequenceParallel (source):

• Parameters (e.g., norm.weight) are marked as Replicate() across TP ranks
• Input is Shard(1) (sharded along the sequence dimension)

2. The Backward Strategy for RMSNorm

The critical logic is in _common_norm_backward_strategy() (source):

# d_weight = sum(grad_out * input / rstd, outer_dim, keepdim=False)
inp_placements = _replicate_dims_start_at(input_src_spec.placements, axis)
reduce_dims_map = _infer_reduce_dims_map(outer_dims, input_src_spec.ndim, False)
out_placements = map_placements_after_reduction(
    inp_placements, outer_dims, reduce_dims_map, "sum"
)

Here's what happens:

Step	Description
1	Input is Shard(1) on the sequence dimension (an "outer dim" for normalization)
2	d_weight requires summing over outer_dims (batch + sequence dimensions)
3	map_placements_after_reduction encounters Shard(dim) where dim is in reduction dimensions
4	The placement is converted to Partial("sum")

Result: d_weight is marked as a Partial DTensor — each rank holds a partial sum of the gradient.

3. When Does the Reduction Actually Happen?

The gradient reduction is deferred and happens when the gradient is consumed:

For Simple FSDP

In simple_fsdp.py, the ReplicateComputation class handles gradient placements:

class ReplicateComputation(torch.nn.Module):
    def __init__(self, ...):
        self.grad_placements: list[Placement] = [
            _ScaledPartial(reduction_divide_factor=...) 
            if reduction_divide_factor is not None
            else Partial(reduce_op="avg")
        ] * self.device_mesh.ndim

When to_local(grad_placements=self.grad_placements) is called on gradients, the Partial placement triggers the reduction — this is where the all_reduce effectively happens.

For Pure TP (without FSDP)

The reduction would happen when:

• optimizer.step() is called, or
• .to_local() is called on the gradient DTensor

---

Why You Don't See It in the Compiled Graph

Reason	Explanation
DTensor abstraction	The compiled graph operates at the local tensor level. DTensor's Partial placement is metadata, and the actual reduction is triggered by redistribution logic when needed.
Fusion with FSDP communication	In Simple FSDP, the gradient reduction for replicated parameters is fused into FSDP's backward reduce_scatter path.
Lazy redistribution	DTensor doesn't insert reductions immediately. It tracks the Partial placement and only triggers reduction when the tensor is redistributed or converted to a local tensor.

---

Where to Look for the Reduction

If you want to trace where the reduction happens:

1. Look for reduce_scatter_tensor operations in your graph — these handle the combined reduction+sharding for FSDP
2. The reduction for replicated parameters happens implicitly when gradients with Partial placement are converted via to_local(grad_placements=[Partial()]) or redistributed

---

Verification

You can verify this behavior by:

# Add a hook on RMSNorm parameters to inspect gradient DTensor placements after backward
def hook(grad):
    print(f"Gradient placements: {grad.placements}")  # Should show Partial("sum") on TP dim
    return grad

model.layers["0"].attention_norm.weight.register_hook(hook)

• Before reduction: gradient should show Partial("sum") on the TP mesh dimension
• After optimizer step or explicit redistribution: gradient becomes Replicate with the reduced values

[tianyu-l]

@alpemreacar
Thanks for pointing out the issue. I believe this is a bug.

https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/simple_fsdp/simple_fsdp.py#L221
On this line, when we do to_local(), it'll lose the grad placement information that on the TP mesh, Replicate RMSNorm weight would see Shard inputs (on sequence dimension). The direct fix is to change to_local() to to_local(in_grad_placements=(.., Partial(), ...) on the TP mesh.

cc @sanketpurandare This is exactly the issue with calling to_local without caring about grad_placements when we are discussing PP api.

While this can be fixed, I think we should do full DTensor everywhere, at least for the simple fsdp folder. cc @fegin

[wwwjn]

@acisseJZhong will working on a fix by adding gradient placement in to_local()