Skip to main content

Training Memory, Serving Capacity, and MoE

Three first-order models for questions students ask after roofline analysis.

analysis
advanced
Use MLSys·im to explain training memory pressure, size a serving deployment, and model MoE routing imbalance without leaving the analytical framework.

The Question

Once a roofline analysis is done, three follow-up questions appear in real design reviews:

  • Why does training need much more memory than inference?
  • How many replicas does a QPS and P99 target require?
  • How much does MoE hot-expert imbalance change the communication cost?

These questions are not separate systems. They are different views of the same demand–supply framework: tensors consume memory, requests consume serving capacity, and routed tokens consume network bandwidth.

NoteLearning Objectives
  • Break training memory into weights, gradients, optimizer state, activations, and communication buffers.
  • Turn a serving target into a first-pass replica count.
  • Sweep a simple MoE routing imbalance factor and read the all-to-all cost.
  • Decide which numbers are direct byte counts and which require calibration.

1. Training Memory Is Not Inference Memory

Inference mainly stores model weights and KV cache. Training also stores gradients, optimizer state, activations, and communication buffers. That is why a model that fits for inference can fail during training.

from mlsysim import Hardware, Models
from mlsysim.solvers import TrainingMemoryModel
from mlsysim.show import table

memory = TrainingMemoryModel().solve(
    model=Models.Language.Llama3_8B,
    hardware=Hardware.Cloud.H100,
    batch_size=8,
    seq_len=2048,
    dp_size=8,
    zero_stage=2,
)

rows = [
    ["Weights", memory.weights],
    ["Gradients", memory.gradients],
    ["Optimizer", memory.optimizer_state],
    ["Activations", memory.activations],
    ["Buffers", memory.communication_buffers],
    ["Total", memory.total_memory],
]

table(["Component", "Per-GPU Memory"], rows)
print(f"Fits on H100: {memory.feasible}")
Component    Per-GPU Memory
───────────────────────────
Weights            16.06 GB
Gradients           2.01 GB
Optimizer          12.04 GB
Activations         9.13 GB
Buffers            0.100 GB
Total              39.34 GB
Fits on H100: True

The first three terms (16.1 GB of weights, 2.0 GB of gradients, and 12.0 GB of optimizer state per rank under ZeRO-2 across 8 data-parallel ranks) are direct byte counts, bringing total per-GPU demand to 39.3 GB (45.8% of the H100’s 86 GB HBM). Activation memory (9.1 GB) is the term most sensitive to framework behavior because checkpointing changes what the backward pass must store versus recompute.

2. Serving Capacity Combines Three Walls

A serving deployment is not sized by TTFT alone. Sizing also requires base request latency, per-replica token capacity, and queueing pressure under load.

from mlsysim.solvers import ServingCapacityModel

capacity = ServingCapacityModel().solve(
    model=Models.Language.Llama3_8B,
    hardware=Hardware.Cloud.H100,
    qps=20,
    target_p99_latency_ms=2000,
    max_seq_len=1024,
    output_tokens=64,
    max_batch_size=16,
    efficiency=0.35,
)

print(f"Required replicas: {capacity.required_replicas}")
print(f"Capacity:          {capacity.qps_capacity:.1f} QPS")
print(f"Utilization:       {capacity.utilization:.1%}")
print(f"Estimated P99:     {capacity.estimated_p99_latency:~.1f}")
print(f"Bottleneck:        {capacity.bottleneck}")
Required replicas: 1
Capacity:          43.4 QPS
Utilization:       46.0%
Estimated P99:     1317.3 ms
Bottleneck:        Feasible

At 20 QPS with a 2000 ms P99 target, 1 H100 replica suffices (43.4 QPS capacity at 46.0% utilization, combining 1154 ms base request latency and 163 ms queueing wait for an estimated P99 of 1317 ms). The efficiency parameter is exposed because the compute-bound part of serving depends on implementation quality. Use the default for quick comparisons; use a measured efficiency value before making a production SLA commitment.

3. MoE Routing Imbalance

Mixture-of-Experts models reduce compute by activating only a subset of experts, but routed tokens create expert-parallel all-to-all traffic. Perfectly balanced routing is an idealization; hot experts increase the effective active work and the routed payload.

from mlsysim import Systems, ureg
from mlsysim.solvers import MoERoutingModel
from mlsysim.models.types import SparseTransformerWorkload

moe = SparseTransformerWorkload(
    name="Toy-MoE-64B",
    architecture="Sparse Transformer",
    parameters=64e9 * ureg.count,
    active_parameters=8e9 * ureg.count,
    experts=8,
    active_experts_per_token=2,
    layers=32,
    hidden_dim=4096,
    heads=32,
)

rows = []
routing_by_gamma = {}
for gamma in [1.0, 1.25, 1.5, 2.0]:
    routing = MoERoutingModel().solve(
        model=moe,
        batch_size=4,
        seq_len=2048,
        ep_size=8,
        routing_imbalance_factor=gamma,
        fleet=Systems.Clusters.Research_256,
    )
    routing_by_gamma[gamma] = routing
    rows.append([
        gamma,
        f"{routing.effective_active_experts:.2f}",
        routing.token_dispatch_bytes,
        routing.all_to_all_latency,
    ])

table(["Imbalance", "Effective Experts", "Routed Bytes", "All-to-All"], rows)
Imbalance  Effective Experts  Routed Bytes  All-to-All
──────────────────────────────────────────────────────
1.00                    2.00      117.4 MB     9.75 ms
1.25                    2.50      146.8 MB    12.09 ms
1.50                    3.00      176.2 MB    14.44 ms
2.00                    4.00      234.9 MB    19.14 ms

This model does not simulate a router. It provides a clean sensitivity knob: if measured routing logs show a 25% hot-expert effect, setting routing_imbalance_factor=1.25 increases the routed payload from 117.4 MB to 146.8 MB and raises all-to-all communication latency from 9.75 ms to 12.09 ms.

4. What Counts as Validation

For these models, validation means matching the level of the approximation:

Output How to validate
Weight, gradient, optimizer, KV-cache bytes Compare against direct tensor counts.
Activation memory Compare against framework memory traces for one model/config.
Serving replica count Benchmark one deployment point, calibrate efficiency, then sweep.
MoE routing imbalance Measure tokens per expert and feed the observed imbalance into the model.

MLSys·im should be used to identify the binding constraint and compare design options before benchmarking. It should not replace empirical measurement for a final production SLA.

Next Steps

Use these models inside a larger analysis:

  • TrainingMemoryModel before DistributedModel to rule out impossible training configurations.
  • ServingCapacityModel before EconomicsModel to convert traffic targets into fleet size and cost.
  • MoERoutingModel with DistributedModel(..., moe_routing_imbalance_factor=...) to see whether expert parallelism is limited by bandwidth.
Back to top