Bug: FP8 + TP≥2 causes NaN loss from step 3 (use_local_output=False incompatible with torchao amax)

[zzhaobh]

Bug description

## Bug: FP8 + TP≥2 causes NaN loss from step 3 (use_local_output=False incompatible with torchao amax)

### Summary

After torchtitan commit eb518a1d8 changed all `ColwiseParallel` layers to
`use_local_output=False`, training with `dtype=fp8` and `TP≥2` produces NaN
loss starting at step 3, with an explosive grad_norm spike at step 2.

### Versions

### Environment

- PyTorch: 2.11.0
- torchao: (latest)
- GPUs: 8× H800, TP=2, EP=4
- Model: Qwen3-30B-A3B (MoE)
- FP8 recipe: `rowwise`, `fp8_recipe_name="rowwise"`

### Steps to Reproduce

Run any training with:
- `tp > 1`
- `dtype = "fp8"` + `fp8_recipe_name = "rowwise"`
- torchao's `Float8ColwiseParallel` / `Float8RowwiseParallel` TP plan

### Observed Behavior

step: 1 loss: 2.833 grad_norm: 105.0
step: 2 loss: 2.831 grad_norm: 572.0 ← spike
step: 3 loss: nan grad_norm: nan ← permanent NaN
step: 4 loss: nan grad_norm: nan
...

### Root Cause

`torchao`'s `Float8ColwiseParallel` computes the per-tensor amax by calling
`.to_local()` on its output activation before quantizing to FP8.

Before eb518a1d8, `ColwiseParallel` used `use_local_output=True`, so the
output was already a plain `torch.Tensor` (the local shard). `.to_local()` was
effectively a no-op and returned the full local tensor.

After eb518a1d8, `use_local_output=False` is set on all `ColwiseParallel`
layers so that DTensor metadata propagates through the TP region. The output is
now a `DTensor(Shard(-1))`. When `Float8ColwiseParallel` calls `.to_local()` on
a `DTensor`, it gets back only the **local shard** — half the activations for
TP=2. The amax is therefore underestimated by roughly 2×, the FP8 scale is too
small, and activations overflow to `inf`/`NaN` within a few steps.

`Float8RowwiseParallel` is **not affected**: its output is `Replicate()` after
the all-reduce, so `.to_local()` on a `DTensor(Replicate)` returns the full
tensor correctly.

### Minimal Fix (workaround applied on our side)

When `Float8ColwiseParallel` is used, set `use_local_output=True` so the output
is a plain Tensor and amax sees the full local activations:

```python
# FP8 path
colwise_local_output = True   # workaround: Float8ColwiseParallel.amax
                               # calls .to_local(), which returns only half
                               # the activations when output is DTensor(Shard(-1))
# non-FP8 path
colwise_local_output = False  # correct DTensor propagation

The proper fix is for torchao's Float8ColwiseParallel to use
tensor.full_tensor() (or an all-gather) instead of .to_local() when the
output is a DTensor, so the amax reflects the global activation range.

Expected Behavior

FP8 + TP≥2 should train without NaN. The amax should be computed over the
global activation tensor (or at minimum over the full local shard before
sharding), not over a DTensor local view.

[fegin]

Wait, I thought use_local_output only indicates whether to_local() is going to be called or not, it doesn't have any sharding placement implication. The issue may be a real but I'm curious is the root cause is what is described in the context.

[danielvegamyhre]

@zzhaobh could you provide the command and python model config used to repro this?

[zzhaobh]

I use Limo dataset test!

HF_CKPT_PATH = "/mnt/workspace/feynmanzhao/hf_model/Qwen2.5_32B_Instruct/Qwen2.5-32B-Instruct"

def main():
result = Trainer(
model=ModelConfig(
name="qwen2_5/32B-Instruct_joydogy",
dp_shard=8,
dp_replicate=1,
tp=4,
pp=1,
cp=1,
attn_backend="sdpa",
attn_mask_type="causal",
dtype="bf16",
ac_mode="full",
ac_early_stop=True,
),
train=TrainConfig(
local_batch_size=1,
global_batch_size=64,
lr=5e-6,
beta1=0.9,
beta2=0.999,
eps=1e-8,
weight_decay=0.0,
max_grad_norm=0.03,
num_epochs=30,
seq_len=36864,
warmup_ratio=0.1,
min_lr_ratio=0.0,
use_dft_loss=False,
loss_reduction="per_token",
chat_template="qwen2_5",
tokenizer_path=HF_CKPT_PATH,
dump_folder="/mnt/workspace/feynmanzhao/training_outputs/limo_qwen25_sft_hf_aligned_v7",
ckpt_interval=15,
log_interval=1,
enable_wandb=True,
wandb_project="limo-sft",
wandb_run_name="qwen25-32b-limo-hf-aligned-v2-tp4-dp8-nopack",
wandb_mode="offline",
shuffle=True,
num_workers=4,
pin_memory=True,
persistent_workers=True,
),
streams=[
StreamConfig(
local="data/limo_qwen25/mds_sft/train",
data_type="messages",
pack_strategy="no_pack",
epoch_size_margin=2.0,
)
],
hf_ckpt_path=HF_CKPT_PATH,
)