One codebase, six phases, four model families: training and serving Llama, Qwen, Mistral, and multimodal models at scale

Field notes on llama-scale-infra -- the project built earlier in this series -- written up the same way as the rest of this collection: intuition first, then the math traced through a real number, then the actual code, then what breaks and why, then how far the same design travels. The project itself used two configs (Llama-2-7B and Llama-3.1-405B) to prove one point: the model’s code never changes across six phases of scale, only the infrastructure wrapped around it does. This document explains that codebase in detail, then spends its second half testing that claim against three real, different model families -- Qwen, Mistral/Mixtral, and multimodal vision-language models -- to see exactly where “just change the config” stops being true and real code has to change instead.

00 · The one design decision everything else follows from#

❓ Why this framing. Every phase of llama-scale-infra -- single GPU, DDP, FSDP, 3D parallelism, KV-cache-aware serving, observability -- wraps a different amount of infrastructure around the exact same LlamaStyleModel class from common/model.py. That’s not a simplification for the sake of a tidy demo; it’s the actual production discipline real frontier labs follow, and it’s the reason this whole codebase can be described as “one thing” instead of six unrelated projects.

$$\text{what changes, phase to phase} = \text{how the model's state is stored (sharding)} + \text{how its compute is split (parallelism)} + \text{how it survives failure (checkpointing, monitoring)}$$ $$\text{what never changes} = \text{the model's forward pass itself}$$

Keeping this separation explicit is why the second half of this document can ask “what does it take to point this same codebase at Qwen, or Mistral, or a multimodal model” as a real, answerable engineering question rather than a rhetorical one -- the answer is always going to be “how much of the model class has to change,” because the infrastructure around it was built specifically not to care.

01 · The shared foundation: a config-driven model, and the math that plans around it#

❓ Why this topic. Before any phase’s infrastructure makes sense, it’s worth being precise about what common/model.py and common/sizing.py actually compute -- because every later phase’s code is either wrapping the model class from the outside or consuming a number the sizing module produced, never hardcoding either.

The model itself: RoPE, GQA, RMSNorm, SwiGLU, driven entirely by a config object#

LlamaStyleModel takes an ArchitectureConfig -- hidden size, layer count, head counts, vocab size -- and builds the same Llama-family transformer block regardless of whether that config describes a 7B or a 405B model. The mechanism worth walking through once, since every later model family in this document reuses it:

RoPE (rotary position embedding). Rather than adding a learned position vector to the token embedding, RoPE rotates the query and key vectors by an angle that depends on their position in the sequence, so that the dot product between a query at position $m$ and a key at position $n$ depends only on their relative distance $m-n$, not their absolute positions:

$$\text{apply\_rope}(x)_{2i,2i+1} = \big(x_{2i}\cos\theta_i - x_{2i+1}\sin\theta_i,\ \ x_{2i}\sin\theta_i + x_{2i+1}\cos\theta_i\big), \qquad \theta_i = \frac{m}{\text{rope\_theta}^{2i/d}}$$

precompute_rope in common/model.py builds this once, up to max_position_embeddings, and apply_rope slices and applies it fresh for whatever sequence length actually shows up -- the same function serves the 7B config’s 4,096-token window and the 405B config’s 131,072-token window without any change, because rope_theta and max_position_embeddings are just numbers pulled from the YAML.

GQA (grouped-query attention), the part that changes most across model families later in this document. GroupedQueryAttention projects queries into num_attention_heads groups but keys/values into a smaller num_key_value_heads, then repeats each KV head across its group before the attention matmul:

k = k.repeat_interleave(self.group_size, dim=1)
v = v.repeat_interleave(self.group_size, dim=1)

Worked example, the two real configs. Llama-2-7B’s config sets num_key_value_heads = num_attention_heads = 32 -- plain multi-head attention, group_size = 1, the repeat_interleave above is a no-op. Llama-3.1-405B sets num_key_value_heads = 8 against num_attention_heads = 128 -- group_size = 16, so every 16 query heads share one KV head. This single config difference is what common/sizing.py’s max_tp_degree_from_kv_heads later turns into a hard serving-time constraint (chapter 5): the 405B config can only be tensor-parallel-sharded 8 ways at most, the 7B config 32 ways, because you cannot split one KV head’s storage across two GPUs without the serving engine needing to do something a lot more complicated than a clean shard.

The sizing module: turning “will this fit” from a guess into an answer#

common/sizing.py’s zero_stage_memory_per_gpu implements the exact $16N$-bytes-per-parameter formula from this project’s companion notes, split by ZeRO stage:

$$\text{ZeRO-1: } \frac{12N}{P} + 4N \qquad \text{ZeRO-2: } \frac{12N}{P} + 2N \qquad \text{ZeRO-3: } \frac{16N}{P}$$

Worked example, run for real. python -m common.sizing against llama_7b (8 GPUs, one node) reports ZeRO stage 1 fits at 37.1 GB/GPU -- comfortably under an 80GB budget, no tensor or pipeline parallelism needed at all. The same command against llama_405b (2,048 nodes × 8 GPUs = 16,384 GPUs) is where this document’s real narrative starts.

The bug this surfaced, and why it’s worth retelling here. The first version of plan_parallelism only checked memory feasibility, and it reported that pure data-parallel ZeRO-3 training fit across all 16,384 GPUs -- technically true (each shard is tiny), and operationally useless, because a 16,384-way all-gather/reduce-scatter every single layer is not something any real interconnect handles efficiently. The fix -- MAX_COMM_EFFICIENT_ZERO3_DP_DEGREE, a deliberately conservative cap on how large a single ZeRO data-parallel group is allowed to be before the planner requires tensor and pipeline parallelism instead -- is the concrete, code-level version of a fact stated more abstractly in this project’s companion notes: memory feasibility and communication feasibility are two different constraints, and a planner that only checks one will confidently give you the wrong answer. After the fix, the same command produces a real, usable plan: TP=8, PP=4, DP=512, ZeRO stage 1, 50.9 GB/GPU -- tensor parallelism confined to one node (the NVLink-boundary rule), pipeline parallelism added specifically to shrink the leftover data-parallel group back under the communication-efficient bound, not because memory demanded it.

Think of common/sizing.py as the difference between a mover asking “will these boxes fit in the truck” and a mover asking “will these boxes fit in the truck, and can the two of us actually carry them up the stairs in a reasonable number of trips.” The first question alone will cheerfully tell you a truck is fine when the real bottleneck was never the truck’s volume.

📚 Further reading

• This series’ companion notes: distributed-training-frameworks-sharding-field-notes (ch. 1, 3) -- the full ZeRO/FSDP memory derivation this module implements
• Su et al., 2021 -- RoFormer: Enhanced Transformer with Rotary Position Embedding
• Ainslie et al., 2023 -- GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

02 · Phases 1–2: single GPU, naive multi-GPU, and real DDP#

❓ Why this topic. These two phases exist to establish, in order: does the model work at all (phase 1), and what’s the wrong way to add more GPUs before the right way (phase 1b, deliberately included as a cautionary step, then phase 2’s DDP fix).

Phase 1: the correctness baseline#

train_single_gpu.py is deliberately unglamorous -- a plain training loop, a fake random-token batch standing in for a real dataloader, gradient clipping, a cosine LR schedule computed by lr_at_step. Its entire purpose is to be the thing every later phase’s output gets compared against: if FSDP (phase 3) or 3D parallelism (phase 4) ever produces a meaningfully different loss trajectory than this script does on the same data and seed, that’s a bug in the parallelism code, not a modeling difference -- the exact parity-check discipline this project’s companion notes describe in detail.

serve_simple.py is the same idea for serving: one FastAPI process, one model instance, a /generate endpoint that loops token-by-token with no persistent KV cache across requests. It’s slow and it doesn’t batch -- and every later serving phase is a direct evolution of this file’s shape (load model at startup, expose an endpoint, return tokens), not a rewrite from scratch.

Phase 1b: naive multi-GPU, kept specifically to be a bad example#

train_naive_multigpu.py wraps the model in nn.DataParallel -- a single line of code, model = nn.DataParallel(model) -- and that’s deliberate: the entire value of this file is watching what happens next, not reading the wrapper itself. DataParallel is single-process, multi-thread: it replicates the model to every GPU on every forward call, splits the batch, and gathers every GPU’s output back to GPU 0 for the loss and backward pass. The script’s own logging makes the cost visible directly -- printing each GPU’s allocated memory after every step shows GPU 0 sitting measurably higher than the rest, because it alone holds the concentrated loss/backward computation. That lopsidedness isn’t a bug to fix in this file; it’s the reason DDP exists at all.

Phase 2: DDP, and the accumulation detail that’s easy to get wrong#

train_ddp.py fixes the asymmetry with a structural change, not a bigger GPU: multi-process instead of multi-thread -- one Python process per GPU, each computing its own forward, loss, and backward independently, with only the gradients synchronized afterward via an all-reduce.

$$\bar{g} = \frac{1}{P}\sum_{i=1}^{P} g_i \qquad \text{(every rank ends this step holding the identical averaged gradient)}$$

The detail worth reading twice in this file is the gradient-accumulation loop:

sync_context = contextlib.nullcontext() if is_last_microbatch else model.no_sync()
with sync_context:
    logits, _ = model(input_ids)
    loss = model.module.loss_fn(logits, targets) / args.grad_accum_steps
    loss.backward()

Without no_sync(), DDP triggers a full gradient all-reduce on every backward() call -- including the three microbatches in a four-step accumulation window that haven’t finished accumulating yet, wasting three-quarters of the communication for no benefit, since only the gradient state after the final microbatch is actually meaningful to synchronize. no_sync() suppresses that all-reduce on every microbatch except the last, so the communication cost matches the number of optimizer steps taken, not the number of forward/backward passes run underneath them -- a detail so easy to miss that “DDP with gradient accumulation is slower than it should be” is one of the most common real performance bugs in distributed training code, precisely because the naive version still produces a correct loss curve, just a wastefully slow one.

serve_dataparallel.py is the serving-side version of the same “more replicas, not more sharding” idea: it launches $N$ full copies of phase 1’s server, each pinned to its own GPU, behind a least-loaded router that tracks in-flight request counts per replica and always sends a new request to whichever replica currently has the fewest. This is deliberately the first lever for more serving capacity -- scale replicas before reaching for tensor parallelism (phase 5) -- because replication adds no communication cost at all, while tensor parallelism adds a real one on every request’s critical path.

📚 Further reading

• This series’ companion notes: distributed-training-frameworks-sharding-field-notes, ch. 1–2 -- the full DDP/no_sync mechanism and communication-primitive definitions

03 · Phase 3: FSDP/ZeRO and resharding-aware checkpointing#

❓ Why this topic. This is the first phase where llama_405b stops OOMing -- and the two details that make that true (the meta device trick and resharding-aware checkpoint metadata) are exactly the ones a naive FSDP port would miss.

The meta-device trick: sharding a model too big to construct in the first place#

train_fsdp.py doesn’t just call fully_shard(model) on an ordinarily-constructed model -- it builds the model on PyTorch’s meta device first:

with torch.device("meta"):
    model = LlamaStyleModel(cfg.architecture)
model = apply_fsdp(model, mesh)
model.to_empty(device=device)   # now allocate each rank's OWN shard only

Why this matters specifically at 405B scale. Constructing LlamaStyleModel(cfg.architecture) normally allocates real memory for every parameter as each nn.Linear/nn.Embedding is created -- for a 405B-parameter model in bf16, that’s roughly 810GB, which doesn’t fit on any single GPU even transiently, and doesn’t even fit most single machines’ combined GPU memory at once. Building on the meta device instead creates every parameter as a shape-only placeholder with no actual storage; apply_fsdp then decides the sharding plan against those placeholders, and to_empty() is the point where each rank finally allocates real memory -- but only for the shard it actually owns. Skipping this and constructing normally is one of the most common ways a first FSDP port of a very large model fails before training ever starts, with an OOM that has nothing to do with the training memory math and everything to do with the brief, transient full-model construction step nobody meant to pay for.

apply_fsdp itself shards at transformer-block granularity -- a fully_shard() call per block, plus one more for the top-level embeddings/final-norm/lm-head -- rather than one call around the whole model, specifically so each block’s parameters can be gathered and freed independently, letting FSDP’s forward-prefetch overlap the next block’s all-gather with the current block’s compute (this project’s companion notes, ch. 4).

Resharding: the checkpoint problem a naive save/load doesn’t solve#

checkpoint_manager.py’s CheckpointManager wraps torch.distributed.checkpoint rather than a plain torch.save, for one specific reason: a checkpoint saved while training on, say, 64 GPUs needs to remain loadable if the next run’s cluster allocation only has 32, or 128. A naive per-rank save (rank_7_shard.pt) has no way to know what “rank 7’s piece” means under a different world size. dcp.save’s metadata instead records each shard’s global offset within the full, logical tensor -- not “GPU 7’s data” but “elements 4096–8191 of this 32768-element parameter” -- so dcp.load can compute, under whatever the current process group’s shard boundaries happen to be, which of the originally-saved shards overlap each new boundary and read exactly those pieces. This is what makes CheckpointManager.load correct regardless of whether it’s called from the same job that wrote the checkpoint or a differently-sized one resuming after a cluster reallocation.

The register_preemption_handler method closes the loop with this project’s companion notes’ real incident numbers: a SIGTERM (what Slurm and Kubernetes send before preempting a job) triggers an immediate checkpoint write before the process exits, so a preemption costs seconds of lost work rather than however long it had been since the last periodic checkpoint -- the concrete mechanism behind Llama 3’s own reported >90% GPU utilization despite roughly one interruption every three hours during its real 405B training run.

📚 Further reading

• This series’ companion notes: distributed-training-frameworks-sharding-field-notes, ch. 3, 5 -- resharding metadata and checkpoint-on-signal in full depth

04 · Phase 4: tensor and pipeline parallelism, computed not guessed#

❓ Why this topic. By this phase, FSDP alone (phase 3) still isn’t the right answer for llama_405b at real cluster scale -- not because it doesn’t fit in memory (it does, per chapter 1’s revised planner), but because a communication-efficient plan needs tensor and pipeline parallelism composed with the ZeRO sharding from phase 3, not instead of it.

parallel_plan.py: one script, so two files can’t silently disagree#

The design choice worth calling out explicitly: ds_3d_config.json and launch_3d_parallel.sh never hardcode a tensor-parallel or pipeline-parallel degree. launch_3d_parallel.sh instead calls parallel_plan.py --emit_json and parses its output before constructing the actual launch command:

PLAN=$(python -m phase4_tensor_pipeline_parallel.parallel_plan --config "${CONFIG}" --emit_json)
TP=$(echo "${PLAN}" | python -c "import sys, json; print(json.load(sys.stdin)['tensor_parallel_size'])")
PP=$(echo "${PLAN}" | python -c "import sys, json; print(json.load(sys.stdin)['pipeline_parallel_size'])")

This is a small thing that prevents a real, common production bug: a config file and a launch script that were both updated by hand for one cluster shape, and only one of them updated the next time the cluster shape changed. Making the plan itself the single source of truth means changing configs/llama_405b.yaml’s node count automatically changes what gets launched, with no second file to remember to edit.

The actual computed plan, and what each number is paying for#

Running this against the real llama_405b.yaml cluster shape (2,048 nodes × 8 GPUs = 16,384 GPUs) produces TP=8, PP=4, DP=512, ZeRO stage 1. Reading that left to right, as the planner itself derived it:

• TP=8: the largest tensor-parallel degree that fits within one node’s 8-GPU NVLink domain -- capped there deliberately, since tensor parallelism’s frequent, small all-reduces are only fast enough to be worth it at NVLink’s roughly 900GB/s, not across the slower inter-node network.
• PP=4: added specifically because TP=8 alone left a data-parallel group of 16384 / 8 = 2048 GPUs, still far above the communication-efficient bound from chapter 1’s fix -- pipeline parallelism’s job here is entirely to shrink that leftover DP group, not to save memory (ZeRO-1 was already sufficient once TP alone was applied).
• DP=512: 16384 / (8 × 4) = 512, the resulting data-parallel group size, now within the bound the planner enforces.
• ZeRO stage 1: the cheapest ZeRO stage (only optimizer-state sharding, the least communication-hungry) that still fits in memory once TP and PP have already done most of the work -- the planner tries stages in cost order and stops at the first one that fits, exactly the discipline of not reaching for a more complex tool than the numbers require.

ds_3d_config.json embeds a comment making this traceability explicit rather than presenting the TP/PP numbers as if they were chosen by inspection: “these three numbers are exactly what phase4a’s parallel_plan.py computes… regenerate this config from that script’s output rather than hand-editing these numbers for a different cluster size.”

A 3D-parallel launch command with hand-picked TP/PP numbers is like a recipe that says “add the right amount of salt” -- it might be right the day it was written, for the kitchen it was written in, and silently wrong the next time either changes. Deriving the numbers from a script that takes the actual cluster shape as input is the difference between a recipe and a scale.

📚 Further reading

• This series’ companion notes: distributed-training-frameworks-sharding-field-notes, ch. 1–2, 4 -- the full 3D-parallelism decomposition and DeepSpeed’s own PipelineModule/AutoTP

05 · Phase 5: KV-cache-aware serving#

❓ Why this topic. Training infrastructure (phases 1–4) and serving infrastructure share vocabulary (tensor parallelism, sharding) but not constraints -- the GQA head count that barely mattered for training becomes a hard ceiling here, and the memory budget that was “will the optimizer state fit” becomes “how many concurrent users can this deployment actually serve.”

The GQA-head ceiling on tensor-parallel serving#

max_tp_degree_from_kv_heads encodes a constraint training-time tensor parallelism doesn’t share: serving-time TP degree cannot exceed num_key_value_heads, because splitting a single KV head’s storage across two GPUs isn’t something serving engines support cleanly. For llama_7b (32 KV heads, plain MHA), that’s effectively no constraint at all -- TP could go as high as 32 if node size allowed it. For llama_405b (8 KV heads, GQA), TP is capped at 8 no matter how many GPUs are available, which happens to line up with one node’s size here but wouldn’t if the node were larger -- a real, independent reason (distinct from any memory argument) that GQA head count is one of the first numbers worth checking before assuming a bigger TP degree is available as an option at serving time.

The wall this planner actually hit, and the real fix#

Running kv_cache_planner.py against llama_405b at TP=8 in bf16 doesn’t produce a small number and a warning -- it hits a hard wall: $405\text{B} \times 2\text{ bytes} / 8 = 101.2$ GB per GPU for weights alone, more than an 80GB GPU has before a single token of KV cache is counted. This is not an edge case the planner mishandled; it’s the same real constraint actual Llama 3.1 405B serving deployments hit, which is why no serious deployment of a model this size runs pure bf16 weights at this TP degree. The planner’s fallback chain mirrors what real production serving actually does about it, in order:

1. Try bf16 at the KV-head-capped TP degree. For llama_7b, this succeeds trivially (1.7GB/GPU for weights, 66.3GB/GPU left over for KV cache -- room for roughly 123 concurrent 8,192-token sequences).
2. Fall back to FP8 weights (and FP8 KV cache) at the same TP degree, halving both footprints. For llama_405b, this is what actually works: weights drop to 50.6GB/GPU, leaving 17.4GB/GPU for KV cache -- enough for roughly 65 concurrent 8,192-token sequences.
3. Only if FP8 still doesn’t fit, add pipeline-parallel inference across nodes -- since PP has no KV-head ceiling the way TP does, it’s the correct next lever once TP is maxed out, not a sign something went wrong.

notes.append(
    f"bf16 weights alone at TP={tp} (the KV-head-capped ceiling) exceed the {gpu_memory_gb}GB "
    f"budget -- this is the same real wall Llama 3.1 405B's own serving deployments hit; "
    f"the production answer is quantized weights, not more bf16 GPUs alone. Trying FP8."
)

serve_vllm.py consumes exactly this computed plan -- tensor_parallel_size, pipeline_parallel_size, and quantization="fp8" if the plan called for it -- rather than a human picking flags by hand, and batching_router.py generalizes phase 2’s simple in-flight-count router with real health checks (a replica that fails its probe is pulled from rotation, not silently sent traffic that will time out) and a queue-depth signal usable directly for Little’s-Law-driven autoscaling.

📚 Further reading

• This series’ companion notes: the long-context/inference companion document, ch. 2, 8 -- PagedAttention, FP8 KV cache, and real vLLM serving code in full depth

06 · Phase 6: observability and reliability#

❓ Why this topic. Everything in phases 1–5 is correct and efficient the day it’s written. This phase is about what happens to that correctness and efficiency over the following three weeks of continuous, unattended operation -- where this project’s companion notes’ real incident data (OPT-175B’s roughly 105 restarts, Llama 3’s ~419 interruptions over 54 days) is the standing reminder that “unattended” doesn’t mean “uneventful.”

experiment_tracking.py’s ExperimentTracker logs more than loss on purpose -- grad norm (which precedes most loss spikes), and MFU, computed directly from the model’s own FLOPs-per-token and measured tokens/sec:

def mfu(self, tokens_per_second: float, gpu_peak_flops: float = 989e12) -> float:
    achieved_flops = self._model_flops_per_token * tokens_per_second
    return achieved_flops / gpu_peak_flops

MFU is the one number that tells you whether phase 4’s carefully-computed 3D-parallelism plan is actually paying off in practice, rather than just being correct on paper -- a plan with the right TP/PP/DP numbers can still run at low MFU if communication/compute overlap (this project’s companion notes, ch. 4) isn’t actually happening, and the tracker logs a warning below 20% specifically to make that gap visible during the run instead of discovered afterward.

monitoring.py’s ThroughputMonitor is a distinct detector from a crash: it catches “nothing threw an exception, but this run is quietly 15% slower than it was yesterday” by comparing a rolling window of recent throughput against a stored baseline -- the failure mode that a system only watching for exceptions never sees at all.

failure_handling.py’s IncidentSeverity levels exist to solve a specific organizational problem, not just a technical one: a retry that succeeds on its second attempt should log quietly (AUTO_RECOVERED), not page a human at 3am for a transient failure the system already fixed itself. Treating every anomaly as equally urgent is what trains an on-call rotation to start ignoring alerts -- a real reliability failure mode that has nothing to do with GPUs or parallelism and everything to do with how a paging system is designed.

📚 Further reading

• This series’ companion notes: distributed-training-frameworks-sharding-field-notes, ch. 5–6 -- the full real-incident data and debugging/observability toolkit this phase’s code implements

07 · Same code, different config: Qwen#

❓ Why this topic. Qwen2.5-7B is architecturally close enough to Llama that it’s a genuine test of the config/code boundary this project draws -- and it fails that test in exactly one specific, instructive place.

What’s config-only#

Qwen2.5-7B’s real published numbers -- 28 layers, hidden size 3584, 28 query heads with 4 KV heads (GQA), head dimension 128, intermediate size 18,944, vocabulary size around 151,646–152,064, rope_theta of 1,000,000, context length up to 131,072 tokens -- are every one of them fields ArchitectureConfig already has. A configs/qwen2_5_7b.yaml with these values, run through the exact same LlamaStyleModel, common/sizing.py, and every phase’s training/serving scripts, would produce a correct parallelism and memory plan with zero code changes: at 4 KV heads, max_tp_degree_from_kv_heads would cap serving-time tensor parallelism at 4 (tighter than Llama-405B’s 8, because Qwen2.5-7B’s GQA ratio is more aggressive at this size), and plan_parallelism would almost certainly conclude -- correctly -- that this model needs no tensor or pipeline parallelism at all, fitting comfortably on a single node exactly the way llama_7b does.

What actually needs a code change: QKV bias#

Qwen2’s attention layers use a bias term on the Q, K, and V projections -- a specific, documented architectural choice distinguishing it from Llama’s bias-free projections. GroupedQueryAttention in common/model.py hardcodes bias=False:

self.q_proj = nn.Linear(cfg.hidden_size, self.n_heads * self.head_dim, bias=False)

Pointing a Qwen config at this class as-is would silently build the wrong architecture -- it would run without error and produce numbers, just not Qwen’s actual numbers, because a whole set of learnable bias parameters the real model depends on simply wouldn’t exist. The fix is small (add a qkv_bias: bool field to ArchitectureConfig, thread it into the three nn.Linear calls) but it’s a genuine code change, not a config value -- the first real data point in this document for where the config/code line actually sits: attention-pattern details that change which parameters exist need code; details that only change how many of an existing parameter type there are (heads, layers, hidden size) don’t.

What’s out of scope, honestly#

Qwen’s real tokenizer is a byte-level BPE vocabulary trained specifically for strong multilingual (especially Chinese/English) compression -- this project’s fake random-token dataset never implements a real tokenizer at all for any config, so “swap in Qwen’s tokenizer” is a real, separate piece of work this codebase doesn’t currently touch, worth naming rather than glossing over.

Qwen’s MoE variants: a bigger change than QKV bias#

Qwen2-MoE and Qwen3-MoE replace the dense FFN with a sparse mixture of experts, the same architectural family as Mixtral (next chapter). This is not a config-only change or even a small code change -- it requires the actual expert-parallel training and serving machinery from this series’ companion notes (distributed-training-frameworks-sharding-field-notes, ch. 2’s expert_parallel_forward), which doesn’t exist anywhere in the current llama-scale-infra codebase at all. Worth flagging plainly: MoE support is the single largest gap between what this project currently implements and what a real Qwen-MoE deployment needs.

08 · Same code, different config: Mistral and Mixtral#

❓ Why this topic. Mistral 7B is architecturally almost identical to Llama at the tensor-shape level -- and its one real difference (sliding window attention) is a cleaner, more clear-cut example of “this needs code, not config” than Qwen’s QKV bias was.

What’s config-only#

Mistral 7B’s real numbers -- 32 layers, hidden size 4096, 32 attention heads, 8 KV heads (GQA), head dimension 128, intermediate size 14,336, vocabulary size 32,000 -- map onto ArchitectureConfig with no new fields needed at all, and are in fact numerically close enough to the llama_7b.yaml config already in this project that a configs/mistral_7b.yaml would look almost identical to it.

What needs a real code change: sliding window attention#

Mistral’s signature architectural choice is that every token only attends to the most recent W tokens (commonly 4,096) rather than the full causal history -- the companion long-context document’s chapter 1 covers why this matters for cost; the point here is what it costs this codebase to actually implement. GroupedQueryAttention.forward currently calls:

out = F.scaled_dot_product_attention(q, k, v, is_causal=kv_cache is None)

is_causal=True gives PyTorch’s fused kernel a full lower-triangular mask -- every token can see every earlier token. Sliding window attention needs a banded mask instead (a token can see the last W tokens and no further back), which is_causal=True alone cannot express. The real fix is either building an explicit band mask and passing it via SDPA’s attn_mask argument, or -- the production-realistic answer -- reaching for flash-attn’s dedicated window_size argument, which implements exactly this pattern with a fused kernel rather than materializing a mask tensor at all. Either way, this is unambiguously a change to common/model.py, not a new YAML field: sliding window attention changes which computation happens, not just how large the existing computation’s tensors are, which is the same category of change QKV bias was for Qwen, just in the attention pattern rather than the projection layers.

Mixtral: a mixture-of-experts change, with real numbers#

Mixtral 8x7B keeps Mistral’s exact attention stack (same 4096/32/8 configuration, same sliding window) and replaces the dense SwiGLU feed-forward block with 8 expert networks and a router selecting the top 2 per token -- 46.7B total parameters, of which only about 12.9B are active for any given token. Implementing this means TransformerBlock’s self.mlp = SwiGLU(cfg) line becomes a router plus a nn.ModuleList of 8 SwiGLU-shaped experts, and -- critically for anything beyond a single-GPU toy run -- training or serving this at real scale means reaching for the exact expert_parallel_forward all-to-all dispatch/combine pattern from this series’ companion notes (distributed-training-frameworks-sharding-field-notes, ch. 2), placing different whole experts on different GPUs rather than trying to fit all 8 experts’ weights on every GPU that holds a copy of the model. This is the same real gap chapter 7 named for Qwen-MoE -- expert parallelism is genuinely new infrastructure this project doesn’t currently include, not a small edit to existing code.

09 · What actually changes for multimodal models#

❓ Why this topic. Multimodal vision-language models are the biggest test of the config/code boundary in this document -- and, honestly, the first case where the answer is “a meaningfully new piece of the model,” not a small patch.

The real architecture, two ways, both published#

Qwen2.5-VL-7B pairs an 0.8B-parameter vision transformer (depth 32, width 1280, patch size 14, local window attention with periodic full-attention layers at fixed depths, native dynamic resolution) with the same 6B-parameter Qwen2.5 LLM backbone described in chapter 7, fusing the two by projecting the ViT’s patch embeddings into the LLM’s hidden dimension and merging spatially adjacent patches (spatial_merge_size=2) before they join the token sequence -- image content becomes, after this projection, just more tokens in the same sequence the LLM already processes.

Llama 3’s own multimodal extension takes a structurally different fusion approach: a 630M-parameter ViT-H/14 encoder feeds cross-attention adapter layers inserted every four transformer blocks (128 query heads, 8 KV heads in the adapter itself) rather than concatenating image tokens directly into the main sequence -- adding roughly 100B extra parameters on top of the 405B text backbone specifically for this fusion mechanism.

What this means for common/model.py, concretely#

Both approaches require a genuinely new component this project’s LlamaStyleModel doesn’t have at all: a vision encoder (a ViT-shaped stack of transformer blocks, self-attention only, no causal masking, operating on image patches instead of text tokens) and a fusion mechanism -- either a projection-and-concatenate step (Qwen2.5-VL’s approach, the simpler of the two to implement) or dedicated cross-attention layers interleaved at fixed depths (Llama 3’s approach, a larger architectural addition). Neither is a config value; both are new classes and a modified TransformerBlock.forward that optionally attends to vision features. This is a categorically bigger change than either Qwen’s QKV bias or Mistral’s sliding window -- those were edits to an existing computation; this is adding a second model that feeds into the first.

What this means for the parallelism and sizing math#

This is the reassuring part: chapters 1–5’s infrastructure code needs comparatively little change, because the vision encoder is small relative to the LLM backbone in every real published example (0.8B vs. 6B for Qwen2.5-VL-7B; 630M vs. 405B for Llama 3’s multimodal variant) -- small enough that it’s a reasonable design choice to keep it data-parallel-replicated (or, in a pipeline-parallel setup, as an early pipeline stage of its own) while the LLM backbone still gets the full FSDP/3D-parallelism treatment from phases 3–4 unchanged. common/sizing.py’s $16N$-bytes-per-parameter formula and plan_parallelism’s TP/PP search still apply directly to the LLM backbone’s parameter count; the vision encoder’s much smaller parameter count would, in essentially every real deployment, never be the thing forcing a bigger TP or PP degree.

What this means for serving and the KV cache#

Once image patches are fused into the sequence (whichever fusion approach is used), they consume KV cache exactly like text tokens -- kv_cache_bytes_per_token and max_tp_degree_from_kv_heads in common/sizing.py don’t need to know a given token came from an image rather than a word; they only need the effective sequence length, which for a multimodal request is text tokens plus however many (merged, projected) patch tokens the image contributed. The practical serving consequence: a single high-resolution image can easily cost as many effective tokens as several paragraphs of text, which is exactly why kv_cache_planner.py’s “max concurrent sequences at a given sequence length” number needs to be computed against the real expected effective sequence length for multimodal traffic, not the text-only assumption this project’s current configs use throughout.

Adding vision to this codebase is less like changing a recipe’s ingredient list (Qwen’s QKV bias, Mistral’s window) and more like adding a whole second kitchen that hands finished dishes through a pass-through window into the first one. The first kitchen (the LLM backbone, and every phase of training/serving infrastructure built around it in this project) barely needs to change its own workflow -- it just needs to accept that some of what comes through the window now started somewhere else.

10 · Closing: the one lesson that generalizes#

Three real model families, three different answers to “what has to change,” and the pattern across all three is the actual takeaway of this whole exercise:

• Qwen: almost entirely config (head counts, hidden size, vocab, rope theta) -- one small, real code change (QKV bias) -- and one large missing piece (MoE expert parallelism, for the MoE variants specifically).
• Mistral/Mixtral: almost entirely config -- one real code change to the attention pattern itself (sliding window, not just its dimensions) -- and, for Mixtral, the same large missing piece (expert parallelism) Qwen-MoE needs.
• Multimodal: the first case where the honest answer is a genuinely new component (a vision encoder plus a fusion mechanism), not a small patch -- but even here, the infrastructure built across phases 1–6 barely changes, because it was written to scale with parameter count and KV-cache footprint, not with any assumption about what kind of tokens produced them.

The dividing line, stated plainly: changes to how many of something exist (layers, heads, hidden dimensions) are config. Changes to which computation happens (an attention mask shape, a bias term, a routing mechanism, an entirely new encoder) are code. Every phase of infrastructure in this project -- sharding, parallelism, checkpointing, serving, observability -- was built to only ever need to know the first kind of number, which is exactly why it kept working, unmodified, across a 60x jump in parameter count from Llama-7B to Llama-405B, and why the same infrastructure needs only small, identifiable, honestly-scoped changes to reach Qwen and Mistral, and one real new component to reach multimodal -- never a rewrite.

---

Zero → Frontier Engineering, log 05. A practical reference, not a tutorial: assumes logs 01–04 and gets denser from there. Named production incidents, papers, and numbers throughout are cited inline per chapter under “further reading.”