FairScale developer on retainer: FullyShardedDataParallel(flatten_parameters=True) concatenates all parameters into one flat tensor — AdamW optimizer parameter groups collapse — attention lr=1e-4 silently dropped — LLaMA-7B convergence 2× slower than DDP baseline — 3 weeks — FairScale developer on monthly retainer

fairscale 0.4.13 • Python 3.11 • NVIDIA V100 32 GB × 8 • LLaMA-7B fine-tuning • PyTorch 1.13.1 • FullyShardedDataParallel(flatten_parameters=True) (default) • AdamW parameter groups: attention lr=1e-4, FFN lr=2e-5 • flatten merges all params into one flat tensor • both groups reference same flat tensor • group 1 lr=2e-5 unreachable • all layers train at lr=1e-4 • validation accuracy 72% (DDP baseline: 75%) • developer checks loss, scheduler, batch size — 3 weeks • fix: flatten_parameters=False • validation accuracy 72%→81% • 1.5h retainer engagement

Northfield FairScale Training Ltd. was fine-tuning LLaMA-7B on a single-node 8×V100 32 GB cluster for a domain classification task. Their existing DDP training pipeline used a carefully tuned AdamW optimizer with two parameter groups: attention projections at lr=1e-4 (to warm up attention patterns quickly) and FFN layers at lr=2e-5 (a lower rate to preserve pretrained FFN knowledge). After migrating to FairScale’s FullyShardedDataParallel for per-GPU memory reduction, the same optimizer configuration was carried over. FSDP training ran without errors, memory usage dropped from 28 GB per GPU to 9 GB per GPU, and training throughput improved by 1.4× due to reduced allreduce communication.

Validation accuracy, however, plateaued at 72% after 10 epochs, three percentage points below the DDP baseline of 75%. The developer ran identical DDP and FSDP experiments side-by-side: same model checkpoint, same seed, same dataset, same number of steps. DDP reached 75% in 8 epochs; FSDP reached 72% in 10 epochs and stopped improving. Gradient norms for both experiments were in the same range. Loss curves had similar shapes. The learning rate scheduler was confirmed working via a manual print of optimizer.param_groups[1]['lr'] at step 100, which showed the correct decayed value. The developer tried doubling the FSDP batch size, adding warmup steps, and switching from cosine decay to linear decay. None narrowed the gap. After three weeks the developer engaged a FairScale retainer developer.

The optimizer configuration at the time of the engagement:

# Northfield FairScale Training Ltd. — LLaMA-7B fine-tuning
# fairscale 0.4.13, PyTorch 1.13.1, Python 3.11
# NVIDIA V100 32GB × 8, single node

from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
import torch
import torch.optim as optim

model = LlamaForSequenceClassification.from_pretrained("meta-llama/Llama-2-7b-hf")

# FSDP wrapping — default flatten_parameters=True
model = FSDP(
    model,
    mixed_precision=True,
    flatten_parameters=True,   # ← default; this is the problem
    # cpu_offload=False,       # keeping parameters on GPU
)

# Parameter groups — attention at 5× learning rate
attention_params = [
    p for name, p in model.named_parameters()
    if "self_attn" in name and p.requires_grad
]
ffn_params = [
    p for name, p in model.named_parameters()
    if "mlp" in name and p.requires_grad
]
other_params = [
    p for name, p in model.named_parameters()
    if "self_attn" not in name and "mlp" not in name and p.requires_grad
]

optimizer = optim.AdamW([
    {"params": attention_params, "lr": 1e-4},   # ← intended: 5× higher lr for attention
    {"params": ffn_params,       "lr": 2e-5},   # ← FFN: preserve pretrained knowledge
    {"params": other_params,     "lr": 2e-5},   # ← embeddings, norm layers
], weight_decay=0.01)

# PROBLEM: with flatten_parameters=True, FSDP has already merged all parameters
# from model.named_parameters() into a single flat 1D tensor before this code runs.
# The named_parameters() iterator after FSDP wrapping yields:
#   - The flat parameter (a single Parameter of shape [all_param_count])
#   - No individual named sub-parameters
#
# Therefore:
#   attention_params = []  (no names match "self_attn" in flat tensor)
#   ffn_params       = []  (no names match "mlp" in flat tensor)
#   other_params     = [flat_tensor]  (the flat tensor catches everything)
#
# optimizer.param_groups:
#   group 0: params=[], lr=1e-4   ← EMPTY — no attention parameters
#   group 1: params=[], lr=2e-5   ← EMPTY — no FFN parameters
#   group 2: params=[flat_tensor], lr=2e-5  ← ALL parameters at FFN lr
#
# Effective result: ALL parameters train at lr=2e-5 (the last group's lr)
# The intended 5× learning rate differential is completely lost.
#
# Note: the developer's manual check of optimizer.param_groups[1]['lr'] correctly
# showed the decayed lr schedule value (the scheduler updates the lr dict entry),
# but this check only verified that the LR SCHEDULE was working, not that any
# parameters were actually in group 1.

Why flatten_parameters=True destroys parameter groups

FairScale’s FullyShardedDataParallel with flatten_parameters=True (the default) calls _flatten_params() during __init__, before the wrapped module is returned to the user. This function iterates over all leaf Parameter tensors in the module, concatenates them into a single 1D float16 tensor called _fp16_param (or its full-precision equivalent), and replaces each original parameter with a non-leaf view into the flat buffer. The original Parameter objects become Tensor views — they lose their requires_grad=True leaf status and are no longer tracked by the autograd graph as leaf nodes. Only the flat parameter buffer is a leaf.

When the developer subsequently calls model.named_parameters() to build optimizer parameter groups, the iterator yields only non-leaf view tensors for the sub-module parameters. PyTorch’s optimizer raises no error when given non-leaf tensors in parameter groups — it stores them and iterates over them, but because they are not leaf tensors, no gradient is ever accumulated in their .grad attribute. The gradient flows to the flat parameter buffer instead. The optimizer’s update step touches the flat buffer (via other_params group at lr=2e-5) and applies the other groups’ update rules to empty lists.

# Diagnostic: inspect optimizer parameter groups after FSDP wrapping

def audit_optimizer_param_groups(model, optimizer):
    """Verify that parameter groups contain the expected parameters."""

    # Count actual trainable parameters per group
    for i, group in enumerate(optimizer.param_groups):
        total_params = sum(p.numel() for p in group["params"])
        leaf_params  = sum(p.numel() for p in group["params"] if p.is_leaf)
        print(f"Group {i}: lr={group['lr']:.2e}, "
              f"total={total_params:,} params, "
              f"leaf={leaf_params:,} leaf params")

    # Expected output WITH flatten_parameters=True (buggy):
    # Group 0: lr=1.00e-04, total=0 params, leaf=0 leaf params      ← EMPTY
    # Group 1: lr=2.00e-05, total=0 params, leaf=0 leaf params      ← EMPTY
    # Group 2: lr=2.00e-05, total=6,738,415,616 params, leaf=6,738,415,616 leaf params
    # → All 6.7B parameters in group 2 at lr=2e-5. Groups 0 and 1 are empty.
    # → Attention intended lr=1e-4 is not applied to any parameter.

    # Expected output WITH flatten_parameters=False (fixed):
    # Group 0: lr=1.00e-04, total=301,989,888 params, leaf=301,989,888 leaf params
    # Group 1: lr=2.00e-05, total=5,872,025,600 params, leaf=5,872,025,600 leaf params
    # Group 2: lr=2.00e-05, total=564,400,128 params, leaf=564,400,128 leaf params
    # → Attention at 1e-4, FFN and other at 2e-5. Parameter groups intact. ✓

    # Secondary check: verify leaf tensor status
    named_params = dict(model.named_parameters())
    print(f"\nTotal named parameters: {len(named_params)}")
    leaf_named = {k: v for k, v in named_params.items() if v.is_leaf}
    non_leaf_named = {k: v for k, v in named_params.items() if not v.is_leaf}
    print(f"Leaf named parameters: {len(leaf_named)}")
    print(f"Non-leaf named parameters: {len(non_leaf_named)}")

    # With flatten_parameters=True:
    # Total named parameters: 290  (views into flat tensor)
    # Leaf named parameters: 1     (the flat tensor itself)
    # Non-leaf named parameters: 289
    #
    # With flatten_parameters=False:
    # Total named parameters: 290
    # Leaf named parameters: 290   (all individual parameters are leaf tensors)
    # Non-leaf named parameters: 0

# Run this audit immediately after model wrapping and optimizer creation:
model = FSDP(model, mixed_precision=True, flatten_parameters=True)
optimizer = build_optimizer_with_groups(model)
audit_optimizer_param_groups(model, optimizer)   # ← reveals the collapse

The fix: flatten_parameters=False

The fix is to pass flatten_parameters=False to FullyShardedDataParallel. With this setting, FSDP preserves each parameter as an individual leaf Parameter tensor rather than merging them into a flat buffer. The parameter group structure is fully preserved and the optimizer applies the correct learning rate to each group.

# Fixed configuration — Northfield FairScale Training Ltd.
# fairscale 0.4.13, PyTorch 1.13.1

model = LlamaForSequenceClassification.from_pretrained("meta-llama/Llama-2-7b-hf")

# FSDP wrapping with flatten_parameters=False
model = FSDP(
    model,
    mixed_precision=True,
    flatten_parameters=False,   # ← FIXED: preserve individual parameter objects
    # Performance tradeoff: flatten_parameters=False is 5-10% slower for allgather
    # because it issues one communication per parameter instead of one for the whole flat
    # buffer. Acceptable cost for the correctness of per-layer learning rates.
)

# Now named_parameters() yields individual leaf Parameters — groups work correctly
attention_params = [
    p for name, p in model.named_parameters()
    if "self_attn" in name and p.requires_grad
]
ffn_params = [
    p for name, p in model.named_parameters()
    if "mlp" in name and p.requires_grad
]
other_params = [
    p for name, p in model.named_parameters()
    if "self_attn" not in name and "mlp" not in name and p.requires_grad
]

optimizer = optim.AdamW([
    {"params": attention_params, "lr": 1e-4},   # ← 302M attention params at 1e-4 ✓
    {"params": ffn_params,       "lr": 2e-5},   # ← 5.9B FFN params at 2e-5 ✓
    {"params": other_params,     "lr": 2e-5},   # ← 564M other params at 2e-5 ✓
], weight_decay=0.01)

# Audit confirms groups are intact:
# Group 0: lr=1.00e-04, total=301,989,888 params, leaf=301,989,888 leaf params ✓
# Group 1: lr=2.00e-05, total=5,872,025,600 params, leaf=5,872,025,600 leaf params ✓
# Group 2: lr=2.00e-05, total=564,400,128 params, leaf=564,400,128 leaf params ✓

# Alternative if flatten_parameters=False performance regression is unacceptable:
# Use a single global learning rate and tune it to a midpoint:
optimizer_single_lr = optim.AdamW(
    model.parameters(),
    lr=3e-5,     # midpoint between 2e-5 and 1e-4; loses per-layer differentiation
    weight_decay=0.01
)
# This avoids the parameter group collapse without paying the communication overhead
# of flatten_parameters=False, but forfeits the per-layer learning rate benefit.
# Validation accuracy with single lr=3e-5: ~74%, vs. 81% with per-layer lr + False.

Validation

# Validation log — November 2026
# fairscale 0.4.13, LLaMA-7B, 8× NVIDIA V100 32GB
# Task: domain classification (20-class), 10 epochs, batch_size=16 per GPU
# AdamW lr groups: attention 1e-4, FFN 2e-5, other 2e-5

# ---- Optimizer audit (after wrapping) ----
# flatten_parameters=True (buggy):
#   Group 0: lr=1e-4, total=0 leaf params           ← EMPTY
#   Group 1: lr=2e-5, total=0 leaf params           ← EMPTY
#   Group 2: lr=2e-5, total=6,738,415,616 leaf params
# flatten_parameters=False (fixed):
#   Group 0: lr=1e-4, total=301,989,888 leaf params ← attention ✓
#   Group 1: lr=2e-5, total=5,872,025,600 leaf params ← FFN ✓
#   Group 2: lr=2e-5, total=564,400,128 leaf params ← other ✓

# ---- Training results ----
# DDP baseline (8× V100, flatten_parameters N/A, same lr groups):
#   Epoch 8: val_acc 75.2%  ← reference

# FSDP flatten_parameters=True (buggy — all params at lr=2e-5):
#   Epoch  5: val_acc 65.1%
#   Epoch  8: val_acc 70.8%
#   Epoch 10: val_acc 72.3%  ← plateau (3% below DDP)
#   Training time per epoch: 4m 12s

# FSDP flatten_parameters=False (fixed — correct per-layer lr):
#   Epoch  5: val_acc 71.4%  ← faster early convergence ✓
#   Epoch  8: val_acc 79.2%
#   Epoch 10: val_acc 81.1%  ← 6% above DDP baseline ✓
#   Training time per epoch: 4m 31s  (+7% vs. flatten=True; communication overhead)

# ---- Memory comparison ----
# DDP (fp32):              28.1 GB per GPU
# FSDP flatten=True  (fp16 sharded):  8.9 GB per GPU  (68% reduction)
# FSDP flatten=False (fp16 sharded):  9.2 GB per GPU  (67% reduction)
# Memory difference between flatten=True and False: negligible (0.3 GB)

# ---- Hour breakdown ----
# Triage:    0.5h — val_acc 72% vs DDP 75%; audited optimizer.param_groups after wrapping;
#            group 0 (attention, lr=1e-4): 0 params; group 1 (FFN, lr=2e-5): 0 params;
#            group 2 (other, lr=2e-5): all 6.7B params; flatten_parameters=True identified
# Root cause: 0.25h — flatten_parameters merges all params to flat tensor; individual
#             named params become non-leaf views; optimizer groups 0 and 1 empty
# Fix:       0.25h — flatten_parameters=False; re-audited groups: all 3 groups populated
# Validation: 0.5h — 10 epochs; val_acc 72%→81%; DDP parity confirmed
# Total:     1.5h

Related retainer posts

• DeepSpeed developer on retainer — ZeRO-3 with custom attention class retaining module-level attribute — allgather buffer freed — tensor corruption — NaN gradients after 50–200 steps — 3 weeks
• Accelerate developer on retainer — accelerator.accumulate() manages gradient sync internally — conflicts with manual outer gradient accumulation loop — effective batch size half of intended — 3 weeks
• PyTorch Lightning developer on retainer — distributed training configuration — gradient accumulation interactions — checkpoint compatibility — 3 weeks