Support the fused SwiGLU override with dist-GEMM TP overlap

[anijain2305]

Stack from ghstack (oldest at bottom):

• -> Support the fused SwiGLU override with dist-GEMM TP overlap #4151
• Add dist GEMM attention and FFN backend #4044

overrides/fused_swiglu fuses the SwiGLU gate and up projections into one w13
parameter and runs a Triton silu(gate) * up. It targets FeedForward.Config
without exact=True, so it also claimed the dist-GEMM FFN config and rewrote it
into its own non-overlapping one. Stacking it on tp_gemm_backend="dist_gemm"
therefore put the TP collectives back outside the GEMMs: the job still trained,
and nothing in the logs said the overlap was gone.

The two optimizations are orthogonal, so this makes them compose.
AllGatherFusedSwiGLU subclasses FusedSwiGLU and overrides only forward
with the overlapping schedule -- the w13 parameter, its state_dict hooks and
the Triton kernel are all inherited -- and the fused_swiglu factory dispatches
on the incoming config so a dist-GEMM FFN maps to it.

Because w13 is a single weight this needs plain AllGatherLinear where the
unfused AllGatherFusedFeedForward needs AllGatherLinearMulti. Worth stating
since it motivated the question: that saves no collective. The multi-weight
version already does one all-gather for both halves, and both do one collective
in forward and two in backward. The only saving is the pair of torch.cat calls
the multi version makes in dgrad. AllGatherLinearMulti therefore stays -- it
is the primitive for any set of column-parallel projections sharing an input.

Adding exact=True to the override was considered instead. It would stop the
override claiming the dist-GEMM config, but then asking for both optimizations
would silently get only one; the dispatch is what makes them compose.

Everything stays in overrides/. An earlier attempt factored the w13 layout
into models/common/feed_forward.py as a shared base so the dist-GEMM FFN could
subclass it, but nothing in core would have used it -- w13 is an overrides-only
concept today -- and it meant changing a FeedForward shared by five models to
serve one override. The override borrows _tp_group_from_context and
_warn_once_unfused from the dist-GEMM module so its TP-off fallback behaves
identically; dependencies only need to run core -> overrides, not both ways.

Two details worth noting for review. The gate/up split needs no contiguous():
folding w13's leading axes yields a GEMM weight whose rows alternate gate, up,
so the output columns alternate too and view(-1, 2).unbind(-1) recovers the
halves as strided views, which the Triton kernel accepts (it is passed each
input's strides). And silu_and_mul_op is called directly rather than through
_fused_silu_and_mul, whose local_map declares the 3D (batch, seq, features)
layout -- on this path the activations are already plain local 2D tensors with
no DTensor to map over.

Test Plan:

New tests: composition and checkpoint-layout coverage in
tests/unit_tests/test_fused_swiglu.py, and a 2-rank numerics test against the
stock FFN in tests/unit_tests/test_dist_gemm.py (CUDA-guarded, since
DTensorTestBase otherwise falls back to CPU/gloo where the Triton op is
unregistered). The composition test asserts on the built module type, not the
config type -- the failure mode is a working-but-unoverlapped FFN, which a
config-type check would not catch.

python -m pytest tests/unit_tests/test_dist_gemm.py \
  tests/unit_tests/test_distributed_linear.py \
  tests/unit_tests/test_fused_swiglu.py \
  tests/unit_tests/test_fused_swiglu_override.py \
  tests/unit_tests/test_inference_moe.py -q

38 passed. End to end on 4 GPUs, confirming the override log reports
AllGatherFusedFeedForward.Config -> AllGatherFusedSwiGLU.Config for all six
layers rather than the non-overlapping FusedSwiGLU.Config:

NGPU=4 ./run_train.sh --module llama3 --config llama3_debugmodel_dist_gemm \
  --parallelism.tensor_parallel_degree 2 \
  --override.imports torchtitan.overrides.fused_swiglu.fused_swiglu --training.steps 5

Also run with --compile.enable, and both unchanged paths (dist-GEMM alone, and
the override alone without dist-GEMM) re-run as regression checks. All train.

Authored with the assistance of an AI agent (Claude Code).

[tianyu-l]

This sounds fine, but I didn't get

Adding exact=True to the override was considered instead. It would stop the
override claiming the dist-GEMM config, but then asking for both optimizations
would silently get only one;

Why could we do exact=True and then have two @override one for FeedForward and the other for AllGatherFusedFeedForward?

[tianyu-l]

maybe I didn't carefully read in last PR, by AllGatherFusedFeedForward sounds not a good name

• not only all-gather, but also reduce scatter
• "fused" all-gather coincide with this file's "fused" swiglu -- and after all our feedforward is swiglu

Probably should rename AllGatherFusedFeedForward to AsyncTPFeedForward / TPOverlappedFeedforward / DistGEMMFeedForward?