Summary
The DDP wrap for finetuning sets find_unused_parameters=True:
# task/train.py:309-310
if is_distributed:
model = DDP(model, device_ids=[rank], find_unused_parameters=True)
This is currently necessary because the finetune model, built from a pretraining checkpoint, retains parameters that still have requires_grad=True but are never exercised by the finetune forward (e.g. vocab-prediction heads, cMIM decoder, encoder readout sub-paths the FFN task head doesn't call). Without the flag, DDP raises:
RuntimeError: Expected to have finished reduction in the prior iteration
before starting a new one ... your module has parameters that were not
used in producing loss.
find_unused_parameters=True makes DDP traverse the autograd graph each backward, mark those params "unused," and skip waiting for their gradients. It works, but it's treating a symptom rather than the cause, and it carries real downsides.
Why this is worth addressing
It masks genuine gradient-flow bugs. With the flag on, if a parameter that should be trained silently stops receiving gradients (a wiring/refactor bug), training continues quietly instead of failing loudly with the DDP error. This removes a useful safety net.
Per-iteration overhead.find_unused_parameters=True adds an autograd-graph traversal on every backward to recompute the used/unused set. For a model where the used set is actually static, this is pure overhead with no benefit.
It signals dead weights in the finetune model. The retained pretrain heads consume memory and live in the optimizer/checkpoint even though they're not used for the downstream task. Beyond the DDP overhead, this bloats saved checkpoints and can confuse downstream tooling.
Suggested fix
Make the used/unused parameter set static and explicit, then turn the flag off:
- When building the finetune model from a pretrain checkpoint, freeze or strip the retained pretraining heads (set
requires_grad=False on, or drop, the params not used by the FFN task head) for each checkpoint type (grover_base / cmim / hybrid). - With every remaining trainable parameter used on every forward, set
find_unused_parameters=False for lower overhead and stricter correctness.
If fully enumerating the unused heads per checkpoint type is not feasible right now, at minimum:
- Add a comment/assertion documenting which params are expected to be unused, and
- Log the set of unused parameters once at startup (DDP can report them) so unexpected additions are visible rather than silently absorbed.
Notes
Found during review of #21 (optional multi-GPU DDP finetuning). Not a correctness bug in the current forward (the used/unused set is identical across ranks, so there's no DDP hang risk) — this is about overhead, checkpoint hygiene, and not masking future gradient bugs.
Summary
The DDP wrap for finetuning sets
find_unused_parameters=True:This is currently necessary because the finetune model, built from a pretraining checkpoint, retains parameters that still have
requires_grad=Truebut are never exercised by the finetune forward (e.g. vocab-prediction heads, cMIM decoder, encoder readout sub-paths the FFN task head doesn't call). Without the flag, DDP raises:find_unused_parameters=Truemakes DDP traverse the autograd graph each backward, mark those params "unused," and skip waiting for their gradients. It works, but it's treating a symptom rather than the cause, and it carries real downsides.Why this is worth addressing
It masks genuine gradient-flow bugs. With the flag on, if a parameter that should be trained silently stops receiving gradients (a wiring/refactor bug), training continues quietly instead of failing loudly with the DDP error. This removes a useful safety net.
Per-iteration overhead.
find_unused_parameters=Trueadds an autograd-graph traversal on every backward to recompute the used/unused set. For a model where the used set is actually static, this is pure overhead with no benefit.It signals dead weights in the finetune model. The retained pretrain heads consume memory and live in the optimizer/checkpoint even though they're not used for the downstream task. Beyond the DDP overhead, this bloats saved checkpoints and can confuse downstream tooling.
Suggested fix
Make the used/unused parameter set static and explicit, then turn the flag off:
requires_grad=Falseon, or drop, the params not used by the FFN task head) for each checkpoint type (grover_base / cmim / hybrid).find_unused_parameters=Falsefor lower overhead and stricter correctness.If fully enumerating the unused heads per checkpoint type is not feasible right now, at minimum:
Notes
Found during review of #21 (optional multi-GPU DDP finetuning). Not a correctness bug in the current forward (the used/unused set is identical across ranks, so there's no DDP hang risk) — this is about overhead, checkpoint hygiene, and not masking future gradient bugs.