Estimate FFN Compute, Memory, and Sharding Communication

Analyze the computational cost, memory traffic, and network communication for a standard two‑layer feed‑forward network (FFN) under different sharding strategies during inference. The forward pass is defined as $$Y = \text{activation}(X W_1) W_2$$, where $$X$$ has dimensions $$[B, H]$$, $$W_1$$ is $$[H, F]$$, and $$W_2$$ is $$[F, H]$$. Count each multiply‑add as 2 FLOPs. Assume each stored element occupies $$s$$ bytes, and ignore the arithmetic of the activation function as…