32K-context fine-tuning OOM on an 8-NPU node
A team is fine-tuning a 7B-parameter transformer on one node with 8 Ascend 910B NPUs, each with 64 GB HBM. The task uses 32K-token sequences, micro-batch 1, and gradient checkpointing. The first attempt uses PyTorch DDP with BF16 parameters and FP32 Adam. Every rank OOMs during the first optimizer step. A rough accounting gives 14 GB for BF16 parameters, 14 GB for gradients, 28 GB for FP32 master weights, and 56 GB for Adam moments, before activations. The team proposes reducing sequence length to 8K, which would violate the product requirement for long-document reasoning. They also consider tensor parallel degree 8 but worry about communication overhead. They ask for a memory-feasible strategy that preserves 32K context.
1. Show why DDP cannot fit and estimate the per-rank state memory under full sharding. Then propose a concrete memory-reduction plan that preserves 32K sequences.
Read the worked answer
DDP replicates the full model and optimizer state on every rank. The state alone is 14 + 14 + 28 + 56 = 112 GB per rank, already far above 64 GB before activations, so OOM is guaranteed. With full sharding across 8 ranks, the same state is about 112/8 = 14 GB per rank, leaving roughly 50 GB for activations, temporary buffers, and fragmentation on each 64 GB NPU.
Preserving 32K requires more than sharding. Use FSDP or an equivalent fully sharded optimizer with parameter, gradient, and optimizer-state sharding; keep gradient checkpointing; use FlashAttention-style fused attention to avoid materializing the full attention matrix; add sequence parallelism so activation memory is split across ranks; and avoid unnecessary FP32 casts outside optimizer state. Start with micro-batch 1 and measure actual activation peaks rather than assuming they fit. If the 50 GB headroom is still insufficient, add CPU offload for optimizer state or use context parallelism. Reducing to 8K is a last resort because it changes the task.
2. Compare tensor parallel degree 8 with FSDP for this node. What would you deploy first, and what communication and profiling evidence would decide whether to change?
Read the worked answer
Tensor parallel degree 8 splits weights and activations within each layer across all 8 NPUs, which directly reduces per-rank activation memory at 32K and keeps collectives inside the node over HCCS. Its cost is an all-reduce or reduce-scatter in every transformer layer, which can hurt with micro-batch 1 because the messages are small and latency-dominated. FSDP full sharding reduces optimizer and parameter memory but all-gathers parameters layer by layer and also introduces per-layer collectives; it does not split activation memory unless sequence parallelism is added.
Deploy FSDP full shard plus activation checkpointing and sequence parallelism first if the measured bottleneck is optimizer state, because it preserves the model structure and is simpler to validate. Profile HCCL time, all-gather latency, per-layer memory, and step time. Switch to or combine with tensor parallel 8 if activation memory remains the limit or if FSDP all-gathers dominate. A hybrid TP=4 with FSDP/sequence parallel over the remaining dimension is a reasonable middle ground if TP=8 communication is too expensive. Decide from profiler traces, not from parameter counts alone.