Katkat
Back

Estimate FFN Compute, Memory, and Sharding Communication

AnthropicSoftware Engineer · Onsite
HardSoftware EngineerOnsiteML System Design

Last updated

Estimate standard FFN FLOPs, parameter and activation bytes, and sharding communication, then compare compute, HBM, and network limits with explicit layout assumptions.

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?

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?