Activation and traffic budgets: count lifetimes as well as shards

A dense activation [M,H] has M×H elements. An even P-way token shard stores M×H/P elements per rank when the placement truly partitions those rows. Multiply by bytes per element for raw payload. Gradients, saved tensors, residuals, softmax statistics, masks, allocator overhead, and alignment are separate. A replicated activation instead retains the full payload. A sum-partial product can also have the full coordinate shape locally despite representing only a contribution, so placement must precede the memory calculation.

Suppose an all-gather materializes the full activation while the original shard stays alive. At that instant the counted payload is full plus shard, not simply full/P. A fused kernel or allocator reuse can change that lifetime, but such a reduction needs implementation evidence. Activation recomputation trades saved state for additional compute and possibly repeated communication. It does not remove every temporary, and its useful granularity depends on which tensors dominate the actual peak.

For an explicitly assumed ring all-gather of an evenly sharded tensor with full size B bytes, each rank sends and receives roughly(P−1)B/P payload bytes. A ring all-reduce of a full same-shaped sum-partial tensor has a different two-phase payload count. Other algorithms, links, protocols, chunking, and overlap change observed behavior. Never attach one per-layer communication number to tensor parallelism without naming forward/backward boundaries and whether sequence or context layouts change the operations.

A latency-bandwidth model adds message-step latency to byte-transfer cost, but effective bandwidth is measured at the relevant boundary. Aggregate send bytes sum contributions across ranks and do not equal time on one shared link. Concurrent compute can overlap some movement, while dependencies can leave another part exposed. Increasing parallel degree reduces some local compute and memory but can make fixed communication costs dominate. Measure representative tensor sizes and topology instead of extrapolating solely from an ideal byte count.

Collect traces that identify which allocations overlap during attention exchange, linear redistribution, and backward. A lower steady activation payload can still leave a larger transient peak if a full gather overlaps several saved buffers. For traffic, record group size, tensor dtype, exact layout transitions, sent/received definition, and whether reported time is total collective duration or the nonoverlapped portion of the iteration. Numerical correctness remains a gate: reducing traffic by omitting required gradient contributions is not an optimization.

The roofline lab is deliberately synthetic and concerns local compute/memory ceilings, not collective links. Reset, set I=4 FLOP/byte and B=80 GB/s, predict ceiling 320 GFLOP/s, then Play or Step. Its fixed compute ceiling is 1,000 GFLOP/s. No input describes network ranks, activation lifetimes, NCCL algorithms, or communication overlap. The lab helps keep byte and operation units explicit, while the worked memory-and-ring budget uses a different clearly stated communication model.