Fix for lora dgrad flatten in _linear_dgrad_matmul

[williammtan]

Restore the frozen-weight dgrad flatten in _linear_dgrad_matmul

Problem

One un-flattened matmul made LoRA 5.3x slower than full fine-tuning.

	full fine-tune (--lora-rank 0)	LoRA, current
actor_train	15.2 s	79.8 s
log_probs (fwd-only control)	3.7 s	4.4 s
actor_train_tflops	5.15	1.03

Under LoRA every base linear is frozen, so every base linear takes LinearWithFrozenWeight.backward, which computes its input gradient through _linear_dgrad_matmul with a bare torch.matmul.

def backward(ctx, grad_output):
    """Backward with frozen weight."""
    (weight,) = ctx.saved_tensors
    grad_input = _linear_dgrad_matmul(grad_output, weight)
    ...


def _linear_dgrad_matmul(grad_output: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
    if _should_compute_dgrad_in_fp32(grad_output, weight):
        return grad_output.float().matmul(weight.float())
    return grad_output.matmul(weight)          # <-- no flattening

That bare matmul is fine for most 3-D inputs and catastrophic for one particular layout. Two tensors of identical shape and dtype, at the Qwen2.5-0.5B output layer (V=151936, H=896, seq 512, micro_batch=1):

a = torch.randn(s, b, V, device="cuda", dtype=torch.bfloat16)                   # already [s, b, V]
c = torch.randn(b, s, V, device="cuda", dtype=torch.bfloat16).transpose(0, 1)   # [b, s, V] viewed as [s, b, V]

a.matmul(w)
c.matmul(w)

	stride	is_contiguous()	dispatches	time
a laid out [s, b, V]	(151936, 151936, 1)	True	aten::mm	2.05 ms
c viewed from [b, s, V]	(151936, 77791232, 1)	True	aten::bmm	1157.56 ms

Same shape, same dtype, both reporting contiguous, 565x apart. The is_contiguous() flag ignores size-1 dims, so it cannot tell these apart; what matmul actually needs is for the leading dims to collapse to one dim by stride, and c's do not.

miles hits case c on every step. It calls the model with labels=None (model.py:313, :505) and computes its own PPO loss, which selects the logits.transpose(0, 1).contiguous() branch in gpt_model.py:848. The backward of that transpose is another transpose, so grad_output arrives at the frozen output layer as exactly the strided view above. The resulting bmm treats sequence as the batch dimension: one matrix-vector product per token, each re-reading the entire weight matrix.

Profiling

Full fine tune

LoRA

Fix

Flatten grad_output to 2-D when dim() > 2, matmul, restore the shape; fp32 branch preserved. At micro_batch = 1 the reshape is a view and copies nothing.

This is a regression restore, not a new optimisation: upstream already guards this. The fork lost it when the backward was extracted into the fork-only _linear_dgrad_matmul to add MEGATRON_TRUE_ON_POLICY_LINEAR_DGRAD_FP32.

Reproducing

1. Run the lora training script: miles/examples/lora/run-qwen2.5-0.5B-megatron-lora.sh, with and without this patch.
2. Record actor_train times and you should find ~88% decrease with the patch.

	full fine-tune (--lora-rank 0)	LoRA, current	LoRA, patched
actor_train	15.2 s	79.8 s	17.0 s
log_probs (fwd-only control)	3.7 s	4.4 s	4.5 s
actor_train_tflops	5.15	1.03	4.59

The patch removes 62.9 s/step, 79% of the LoRA training step. Unpatched LoRA is 5.3x slower than full fine-tuning; patched it is within 12%. log_probs is forward-only and flat across all three, localising the regression to the backward pass.

Ran with 1x GB10, TP=PP=CP=EP=1, --num-rollout 3, 64 sequences/step. Medians of steps 1-2.