Estimate FFN Compute, Memory, and Sharding Communication

Estimate the compute, memory traffic, and network communication of a standard two-layer feed-forward network under different sharding strategies. Use those estimates to discuss compute, memory, and communication bottlenecks. ### Constraints & Assumptions For a self-contained practice model, let `X[B,H]`, `W1[H,F]`, and `W2[F,H]` compute `Y = activation(X W1) W2`. This is a nongated inference FFN. Count a multiply-add as two FLOPs, use s bytes per stored element, and ignore activation arithmetic when it is small relative to the matmuls. State caching, materialization, and collective assumptions rather than treating all byte counts as universal measurements. ### Clarifying Questions Which tensors are replicated or sharded at entry and exit? Are weights resident in cache? Is the intermediate materialized? What hardware compute, HBM bandwidth, network bandwidth, and collective latency are available? ### What a Strong Answer Covers Unsharded FLOPs and bytes, per-rank tensor shapes, data/FFN-dimension/hidden-dimension sharding, reductions, pipeline tradeoffs, and a roofline-style lower bound with its limitations. ### Follow-up Questions Why does sharding the FFN expansion dimension need a reduction of the output? Why can hidden-dimension sharding require a larger intermediate reduction? How do microbatching and collective startup latency change the result?