PyTorch Model Design
Implementation-focused skill for ML model development in PyTorch — architecture design,
profiling, distributed training, and production inference.
How to Use This Skill
Six reference files cover distinct concerns. Read only what the current task needs.
| Reference |
When to Read |
| references/principles.md |
Start here. Engineering principles: computation graph, memory model, precision, training stability, parallelism, inference vs training, multimodal conventions, arithmetic intensity / roofline model, data loading at scale, reproducibility, loss function patterns |
| references/architecture.md |
Implementing a model: attention variants (MHA/GQA/MQA/sliding window), FFN variants (SwiGLU/MoE), position encoding (RoPE), block assembly, initialization, gradient debugging |
| references/profiling-optimization.md |
Making it fast: MFU, FLOPs counting, memory analysis, AMP (bf16/fp16/fp8), FlashAttention, torch.compile, gradient checkpointing, sequence packing, benchmarking |
| references/distributed-training.md |
Multi-GPU/multi-node: DDP, FSDP2, tensor parallel, pipeline parallel, DeviceMesh, 2D/3D parallelism, context parallel, distributed checkpointing |
| references/multimodal-temporal.md |
Multimodal and temporal models: modality tokenization (text/image/audio/video), cross-attention vs interleaved, modality alignment, temporal patterns, causal masking, streaming inference, memory budgeting |
| references/inference-deployment.md |
Production inference: KV cache, quantization (INT8/FP8/GPTQ/AWQ via torchao), speculative decoding, continuous batching, torch.export |
Phase 1: Architecture Design
Decision: Graph Mode
Does the model have data-dependent control flow?
(early exit, variable-length loops, conditional layers)
├─ YES → Eager mode. torch.compile(dynamic=True) may work
│ but test carefully — graph breaks eliminate most benefit.
└─ NO ↓
Research iteration mode (architecture changes frequently)?
├─ YES → Start eager. Compile when architecture stabilizes.
│ Compiled model debugging is significantly harder.
└─ NO ↓
Production / throughput-critical?
└─ YES → torch.compile(mode="max-autotune").
Invest in eliminating graph breaks first.
Decision: Architecture Selection
The standard mapping (sequences → Transformer, images → CNN/ViT, graphs → GNN) is textbook.
Focus on non-obvious choices where defaults are wrong:
| Situation |
Choice |
Why |
| Sequence $> 8$K tokens |
SSM (Mamba) or sliding window attention |
Quadratic attention dominates at scale; sub-quadratic trades small accuracy loss for 5-10× throughput |
| Multi-head attention, memory-constrained |
GQA (Grouped Query Attention) |
$k$ KV heads shared by $h/k$ Q heads — LLaMA 2/3, Gemma, Mistral all use this |
| Large capacity without compute cost |
Sparse MoE, top-$k$ routing |
Linear capacity scaling, sublinear compute — but load balancing is non-trivial |
| Tabular data, $< 10$K samples |
XGBoost/LightGBM, not neural |
Trees learn feature interactions from splitting; neural nets need much more data |
| Small labeled + large unlabeled |
SSL pretrain then fine-tune |
Domain-specific SSL beats ImageNet transfer for specialized domains |
| Strict latency ($< 10$ ms) |
Distill large → small |
Distilled small models consistently beat architectures designed small from scratch |
| Multi-modal inputs |
Separate encoders → shared latent space → task head |
Shared trunk forces all modalities through the same bottleneck |
Read references/architecture.md for implementation patterns.
Phase 2: Profiling and Bottleneck Diagnosis
Run top-to-bottom, stop at the first hit:
1. GPU utilization < 80%?
├─ YES → Data loading bottleneck
│ Fix: num_workers, pin_memory, prefetch_factor,
│ WebDataset for streaming, move preprocessing to GPU
└─ NO ↓
2. Memory OOM or > 90% usage?
├─ YES → Activation/gradient memory bottleneck
│ Fix: gradient checkpointing, bf16 AMP, sequence packing,
│ FSDP2 (shards params + grads + optimizer state)
└─ NO ↓
3. MFU < 40%? (MFU = actual throughput / theoretical peak)
├─ YES → Kernel efficiency bottleneck
│ Fix: torch.compile, FlashAttention, operator fusion,
│ eliminate graph breaks, FP8 on H100+
└─ NO ↓
4. Multi-GPU, communication slow?
├─ YES → Communication bottleneck
│ Fix: overlap compute/comm (FSDP2 default),
│ tune NCCL settings, check NVLink vs PCIe topology
└─ NO → Compute-bound (good). Scale with more GPUs.
Optimization Map
| Problem |
Solution |
Impact |
Reference |
| OOM training |
AMP (bf16) |
~50% memory |
profiling-optimization.md §AMP |
| OOM large activations |
Gradient checkpointing |
~60-70% activation memory, ~30% slower |
profiling-optimization.md §Checkpointing |
| Model too large for 1 GPU |
FSDP2 |
Shards params+grads+optimizer |
distributed-training.md §FSDP2 |
| Need more capacity |
Tensor parallel |
Splits weight matrices across GPUs |
distributed-training.md §TP |
| Very deep model, PP needed |
Pipeline parallel |
Splits layers across GPUs |
distributed-training.md §PP |
| Slow training throughput |
torch.compile |
10-40% speedup typical |
profiling-optimization.md §Compile |
| Slow attention |
FlashAttention-3 |
2-4× speedup, O(1) memory |
profiling-optimization.md §FlashAttention |
| H100+ training |
FP8 mixed precision |
~2× throughput vs bf16 |
profiling-optimization.md §FP8 |
| Slow inference |
INT8/FP8 quantization |
2-4× speedup, 50-75% memory |
inference-deployment.md §Quantization |
| Autoregressive inference |
KV cache |
Removes $O(L)$ recomputation per step |
inference-deployment.md §KVCache |
| Single request latency |
Speculative decoding |
2-3× latency reduction |
inference-deployment.md §Speculative |
| Multi-modal architecture |
Cross-attention + modality tokenization |
— |
multimodal-temporal.md |
Phase 3: Distributed Training Strategy
For models too large for a single GPU, choose the parallelism strategy:
Model fits on 1 GPU?
├─ YES → DDP (data parallel only). FSDP2 not needed.
└─ NO ↓
Model fits with FSDP2 (parameter sharding)?
├─ YES → FSDP2 alone. Covers most cases up to ~70B params on 8×A100.
└─ NO ↓
Need tensor parallel (split individual layers)?
├─ YES → FSDP2 + TP (2D parallelism). Standard for 70B+ models.
│ Requires NVLink — PCIe TP is communication-bound.
└─ NO → Consider Pipeline Parallel for depth-heavy models.
Training extremely long sequences (> 32K)?
└─ YES → Context Parallel (CP) — split sequence across GPUs.
Ring Attention pattern. Add as 4th dimension.
Read references/distributed-training.md for setup code.
Output Format
Architecture Design
## Architecture: [Name]
### Design Decisions
| Decision | Choice | Rationale |
|----------|--------|-----------|
| Graph mode | eager / compiled | [why] |
| Attention variant | MHA / GQA / MLA / Sliding Window | [why] |
| Position encoding | RoPE / ALiBi / NoPE | [why] |
| Normalization | Pre-LN RMSNorm / QK-Norm | [why] |
| FFN | SwiGLU / MoE | [why] |
### Implementation Notes
[Key PyTorch patterns, non-obvious gotchas, custom layers needed]
Profiling Report
## Profile: [Model Name]
### Bottleneck: [diagnosis from the 4-step checklist]
| Component | Params | FLOPs/sample | Peak Memory |
|-----------|--------|-------------|-------------|
### MFU: [X]% (target: >50% for large models)
### Top-3 Optimizations
1. [What] → [expected impact] — [reference section]
2. ...
3. ...
1---2name: torch-model-design3description: PyTorch model design, computation graph analysis, profiling, and optimization — from architecture implementation to production-ready training and inference. Use when the user asks about implementing neural architectures in PyTorch, profiling FLOPs/memory/throughput, optimizing training speed or memory usage, debugging gradient flow or numerical stability, choosing between dynamic and static graphs (eager vs torch.compile), setting up distributed training (FSDP2, tensor parallel, pipeline parallel, 2D/3D parallelism), multimodal or temporal model architecture, or optimizing inference (quantization, KV cache, speculative decoding, continuous batching). Also trigger on: computation graph, autograd, torch.compile, FSDP2, DTensor, tensor parallelism, pipeline parallelism, mixed precision, AMP, bf16, fp8, gradient checkpointing, FlashAttention, model profiling, OOM debugging, CUDA memory, inference optimization, model quantization, torchao, multimodal, cross-attention, vision encoder, temporal model.4---56# PyTorch Model Design78Implementation-focused skill for ML model development in PyTorch — architecture design,9profiling, distributed training, and production inference.1011---1213## How to Use This Skill1415Six reference files cover distinct concerns. Read only what the current task needs.1617| Reference | When to Read |18|-----------|-------------|19| [references/principles.md](references/principles.md) | **Start here.** Engineering principles: computation graph, memory model, precision, training stability, parallelism, inference vs training, multimodal conventions, arithmetic intensity / roofline model, data loading at scale, reproducibility, loss function patterns |20| [references/architecture.md](references/architecture.md) | Implementing a model: attention variants (MHA/GQA/MQA/sliding window), FFN variants (SwiGLU/MoE), position encoding (RoPE), block assembly, initialization, gradient debugging |21| [references/profiling-optimization.md](references/profiling-optimization.md) | Making it fast: MFU, FLOPs counting, memory analysis, AMP (bf16/fp16/fp8), FlashAttention, torch.compile, gradient checkpointing, sequence packing, benchmarking |22| [references/distributed-training.md](references/distributed-training.md) | Multi-GPU/multi-node: DDP, FSDP2, tensor parallel, pipeline parallel, DeviceMesh, 2D/3D parallelism, context parallel, distributed checkpointing |23| [references/multimodal-temporal.md](references/multimodal-temporal.md) | Multimodal and temporal models: modality tokenization (text/image/audio/video), cross-attention vs interleaved, modality alignment, temporal patterns, causal masking, streaming inference, memory budgeting |24| [references/inference-deployment.md](references/inference-deployment.md) | Production inference: KV cache, quantization (INT8/FP8/GPTQ/AWQ via torchao), speculative decoding, continuous batching, torch.export |2526---2728## Phase 1: Architecture Design2930### Decision: Graph Mode3132```33Does the model have data-dependent control flow?34(early exit, variable-length loops, conditional layers)35├─ YES → Eager mode. torch.compile(dynamic=True) may work36│ but test carefully — graph breaks eliminate most benefit.37└─ NO ↓3839Research iteration mode (architecture changes frequently)?40├─ YES → Start eager. Compile when architecture stabilizes.41│ Compiled model debugging is significantly harder.42└─ NO ↓4344Production / throughput-critical?45└─ YES → torch.compile(mode="max-autotune").46 Invest in eliminating graph breaks first.47```4849### Decision: Architecture Selection5051The standard mapping (sequences → Transformer, images → CNN/ViT, graphs → GNN) is textbook.52Focus on **non-obvious choices** where defaults are wrong:5354| Situation | Choice | Why |55|----------|--------|-----|56| Sequence $> 8$K tokens | SSM (Mamba) or sliding window attention | Quadratic attention dominates at scale; sub-quadratic trades small accuracy loss for 5-10× throughput |57| Multi-head attention, memory-constrained | GQA (Grouped Query Attention) | $k$ KV heads shared by $h/k$ Q heads — LLaMA 2/3, Gemma, Mistral all use this |58| Large capacity without compute cost | Sparse MoE, top-$k$ routing | Linear capacity scaling, sublinear compute — but load balancing is non-trivial |59| Tabular data, $< 10$K samples | XGBoost/LightGBM, not neural | Trees learn feature interactions from splitting; neural nets need much more data |60| Small labeled + large unlabeled | SSL pretrain then fine-tune | Domain-specific SSL beats ImageNet transfer for specialized domains |61| Strict latency ($< 10$ ms) | Distill large → small | Distilled small models consistently beat architectures designed small from scratch |62| Multi-modal inputs | Separate encoders → shared latent space → task head | Shared trunk forces all modalities through the same bottleneck |6364Read [references/architecture.md](references/architecture.md) for implementation patterns.6566---6768## Phase 2: Profiling and Bottleneck Diagnosis6970Run top-to-bottom, stop at the first hit:7172```731. GPU utilization < 80%?74 ├─ YES → Data loading bottleneck75 │ Fix: num_workers, pin_memory, prefetch_factor,76 │ WebDataset for streaming, move preprocessing to GPU77 └─ NO ↓78792. Memory OOM or > 90% usage?80 ├─ YES → Activation/gradient memory bottleneck81 │ Fix: gradient checkpointing, bf16 AMP, sequence packing,82 │ FSDP2 (shards params + grads + optimizer state)83 └─ NO ↓84853. MFU < 40%? (MFU = actual throughput / theoretical peak)86 ├─ YES → Kernel efficiency bottleneck87 │ Fix: torch.compile, FlashAttention, operator fusion,88 │ eliminate graph breaks, FP8 on H100+89 └─ NO ↓90914. Multi-GPU, communication slow?92 ├─ YES → Communication bottleneck93 │ Fix: overlap compute/comm (FSDP2 default),94 │ tune NCCL settings, check NVLink vs PCIe topology95 └─ NO → Compute-bound (good). Scale with more GPUs.96```9798### Optimization Map99100| Problem | Solution | Impact | Reference |101|---------|----------|--------|-----------|102| OOM training | AMP (bf16) | ~50% memory | profiling-optimization.md §AMP |103| OOM large activations | Gradient checkpointing | ~60-70% activation memory, ~30% slower | profiling-optimization.md §Checkpointing |104| Model too large for 1 GPU | FSDP2 | Shards params+grads+optimizer | distributed-training.md §FSDP2 |105| Need more capacity | Tensor parallel | Splits weight matrices across GPUs | distributed-training.md §TP |106| Very deep model, PP needed | Pipeline parallel | Splits layers across GPUs | distributed-training.md §PP |107| Slow training throughput | torch.compile | 10-40% speedup typical | profiling-optimization.md §Compile |108| Slow attention | FlashAttention-3 | 2-4× speedup, O(1) memory | profiling-optimization.md §FlashAttention |109| H100+ training | FP8 mixed precision | ~2× throughput vs bf16 | profiling-optimization.md §FP8 |110| Slow inference | INT8/FP8 quantization | 2-4× speedup, 50-75% memory | inference-deployment.md §Quantization |111| Autoregressive inference | KV cache | Removes $O(L)$ recomputation per step | inference-deployment.md §KVCache |112| Single request latency | Speculative decoding | 2-3× latency reduction | inference-deployment.md §Speculative |113| Multi-modal architecture | Cross-attention + modality tokenization | — | multimodal-temporal.md |114115---116117## Phase 3: Distributed Training Strategy118119For models too large for a single GPU, choose the parallelism strategy:120121```122Model fits on 1 GPU?123├─ YES → DDP (data parallel only). FSDP2 not needed.124└─ NO ↓125126Model fits with FSDP2 (parameter sharding)?127├─ YES → FSDP2 alone. Covers most cases up to ~70B params on 8×A100.128└─ NO ↓129130Need tensor parallel (split individual layers)?131├─ YES → FSDP2 + TP (2D parallelism). Standard for 70B+ models.132│ Requires NVLink — PCIe TP is communication-bound.133└─ NO → Consider Pipeline Parallel for depth-heavy models.134135Training extremely long sequences (> 32K)?136└─ YES → Context Parallel (CP) — split sequence across GPUs.137 Ring Attention pattern. Add as 4th dimension.138```139140Read [references/distributed-training.md](references/distributed-training.md) for setup code.141142---143144## Output Format145146### Architecture Design147```148## Architecture: [Name]149150### Design Decisions151| Decision | Choice | Rationale |152|----------|--------|-----------|153| Graph mode | eager / compiled | [why] |154| Attention variant | MHA / GQA / MLA / Sliding Window | [why] |155| Position encoding | RoPE / ALiBi / NoPE | [why] |156| Normalization | Pre-LN RMSNorm / QK-Norm | [why] |157| FFN | SwiGLU / MoE | [why] |158159### Implementation Notes160[Key PyTorch patterns, non-obvious gotchas, custom layers needed]161```162163### Profiling Report164```165## Profile: [Model Name]166167### Bottleneck: [diagnosis from the 4-step checklist]168169| Component | Params | FLOPs/sample | Peak Memory |170|-----------|--------|-------------|-------------|171172### MFU: [X]% (target: >50% for large models)173174### Top-3 Optimizations1751. [What] → [expected impact] — [reference section]1762. ...1773. ...178```
guliqianxun