[Reland] Add dist GEMM attention and FFN backend

[anijain2305]

Reland of #4044 because I did a rookie mistake of merging ghstack opened PRs.

Motivation

Async-TP -- folding the TP all-gather and reduce-scatter into the adjacent GEMM rather than running them beside it -- is currently only reachable in torchtitan through a compiler: either torch.compile's micro-pipelining pass behind --parallelism.enable_async_tensor_parallel, or a GraphTrainer graph pass. That ties a parallelism strategy to a particular compilation path.

This makes it a property of the model config instead, selected by a tp_gemm_backend argument on model_registry exactly as attn_backend already selects the attention kernel. Async-TP then works under the eager trainer with no graph pass involved.

Structure

torchtitan/distributed/linear.py holds three reusable autograd Functions. AllGatherLinear is a column-parallel projection (all-gather the sequence shard, then GEMM) and LinearReduceScatter its row-parallel dual (GEMM, then reduce-scatter). AllGatherLinearMulti serves several column-parallel linears that share one input, so the SwiGLU pair w1/w3 costs one all-gather rather than two; its backward still needs only the same two collectives, because the sum over outputs is expressed as a single concatenated product. None of them is attention-specific; MoE projections could use the same pair.

torchtitan/models/common/dist_gemm.py holds the modules. AllGatherFusedQKVLinear, RowParallelLinear and AllGatherFusedFeedForward are drop-in replacements for the stock QKV, output and SwiGLU projections, keeping the stock parameter layouts so checkpoints still interoperate. RowParallelLinear serves both attention's wo and the FFN's w2.

Selection follows the attn_backend path one for one: model_registry ->
flavor function -> _build_llama3_layers -> make_gqa_config / make_ffn_config, which is where tp_gemm_backend="dist_gemm" picks the fused classes. Only llama3 is threaded; every other model keeps the default and can gain the argument when someone needs it. llama3_debugmodel_dist_gemm is the config-registry entry, and the h100 integration test runs it at TP=2.

No attention or FFN subclass is needed beyond the projections themselves. That required one fix in shared code: GQAttention.forward captured B, L from its input and reused them to fold the attention output, which is only valid when the QKV preserves sequence length. AllGatherFusedQKVLinear gathers the SP shard, so it does not, and under spmd_types -- where a tensor's shape is its local shape -- that silently folded sequence into features. It now derives the fold from the tensor being folded.

spmd_types only

The fused modules take and return plain local tensors, which is the spmd_types contract; the DTensor backends are being deprecated and are rejected outright. Two preconditions are checked at parallelize time, because neither is detectable from inside a module once activations carry no placements: the backend must be spmd_types, and sequence parallelism must be enabled. The fused GEMMs are the SP collectives, so with SP disabled there is nothing to gather and wo would reduce-scatter where it must all-reduce.

Because selection happens at config-construction time, set_gqa_attention_sharding and set_dense_ffn_sharding declare the fused blocks' contracts directly: no boundary all-gather, and a row-parallel output that emits its final Shard(1) rather than the Partial a stock rowwise linear produces. Both branches are transitional and collapse once redistribute collectives move inside the modules generally.

Symmetric memory

One workspace per process group serves every layer, and the symm_mem ops size it themselves: each computes its own requirement and get_symm_mem_workspace grows monotonically, so the cost is a max over layers rather than a sum. Nothing here reserves it up front. Growth re-rendezvouses, which is a collective and is rejected during CUDA graph capture, but capture is always preceded by a real warmup call (graph_trainer/cudagraph.py), and a warmup forward and backward runs every layer's collectives -- so all growth has happened before anything is recorded. An earlier revision pre-reserved from parallelize to make that explicit; it bought one growth instead of a few during warmup, at the cost of duplicating the sizing arithmetic, and is not what makes capture safe.

One consequence worth keeping: both wgrad all-gathers leave return_A at its default. Passing False would select _multimem_all_gather_matmul, which reserves the full gathered buffer rather than a shard -- ranks times more symmetric memory for that one call -- and its heuristic evaluates to K <= 2048 there, a dimension it was not tuned for. Leaving it defaulted also keeps forward and backward on the same schedule.

Every op carves its buffers from offset 0 of that shared workspace, so the regions alias; safety comes from the ops being serialized rather than from disjointness. Each op also brackets itself with barriers, so ranks cannot run ahead of one another. Sequential module forwards and autograd backward are single-stream, which is how models actually run, so this holds today -- but deliberate cross-stream overlap would need distinct workspace offsets, not a barrier.

Known gap

--debug.spmd_typechecking does not pass. The autograd Functions are registered with typecheck_forward declaring their type transitions, which the checker requires, but it then rejects the QKV reshape that splits the TP-sharded feature dim into (n_kv, r_dim, head_dim) -- aligned in practice, since the shard is over n_kv. The stock FusedQKVLinear hits the same limitation and works around it with spmd.local() plus an explicit PartitionSpec, carrying a TODO to remove that once the checker handles aligned splits. Rather than add a second instance of a pattern already marked for deletion, this leaves typechecking unsupported until that lands upstream.

Authored with assistance from Claude Code.

ghstack-source-id: e0015932027e7b6c4f5b424bd47d690daa77c964
Pull-Request: #4044