feat: optional multi-GPU (DDP) finetuning#21
Merged
Conversation
Collaborator
|
I wonder if we should have only 1 entrypoint for finetuning. Currently, we have two |
sveccham
requested changes
Jul 9, 2026
sveccham
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for the DDP implementation. Overall, useful useful feature to have. I have left comments in-line.
sveccham
approved these changes
Jul 13, 2026
sveccham
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for the finetune ddp implementation. I have approved the PR. Should be good to merge after a quick rebase.
Adds a DistributedDataParallel launcher for finetuning, mirroring the existing pretraining DDP path so users can accelerate finetunes across multiple GPUs. - finetune_ddp.py: one-process-per-GPU launcher (mp.spawn, WORLD_SIZE env, NCCL topology auto-config, per-rank determinism), routed through cross_validate -> run_training. - kermt/util/ddp_utils.py: shared ddp_setup + configure_nccl_for_topology so the pretrain and finetune launchers share one bootstrap. - run_training / train (task/train.py): thread rank/world_size; DDP-wrap the model (find_unused_parameters=True, since the finetune model carries params the FFN forward does not exercise); DistributedSampler + set_epoch on the train loader; run validation on the unwrapped module over the full val set on every rank (consistent best-model selection); gate checkpoint saving, logging and test-set evaluation to rank 0; all-reduce the standalone MTLLoss.log_sigma gradient. Also use a distinct train_loader instead of reassigning train_data, which fixes the ensemble_size>1 DataLoader double-wrap. - cross_validate: thread rank/world_size; gate wandb + reporting to rank 0. - build_lr_scheduler: world-size-aware steps_per_epoch. Single-process finetune (python main.py finetune ...) is unchanged: DDP only activates when a process group is initialized (i.e. via finetune_ddp.py). Batch size is per-GPU; effective global batch = batch_size * world_size. Signed-off-by: Eva Xue <evax@nvidia.com>
Adds a --num-gpus option to the finetune runner and documents it in the skill. Default is 1 (single-process `main.py finetune`, unchanged). With --num-gpus N (N>1) the runner launches finetune_ddp.py with WORLD_SIZE=N instead, for data-parallel training; --batch-size is per-GPU. - agent/scripts/run_finetune_local.py: --num-gpus arg; _build_argv selects main.py vs finetune_ddp.py; set WORLD_SIZE in the exec + replay env; skip the single-device --gpu flag when distributed (each rank pins its own device); record num_gpus/distributed in run.json. - agent/skills/kermt-finetune/SKILL.md: document --num-gpus across the hardware, arguments, invocation-template, and error-hint sections. Signed-off-by: Eva Xue <evax@nvidia.com>
Address review comments about drift-prone duplicated setup code. Determinism: the identical 6-line block in main.py:setup, main_hpo.py:setup, and finetune_ddp.py:_setup_determinism is hoisted into kermt/util/utils.py as setup_determinism(seed), with seed_rngs(seed) for the RNG-only reset that cross_validate.py uses per fold. setup_determinism now also sets CUBLAS_WORKSPACE_CONFIG (which torch.use_deterministic_algorithms requires for cuBLAS on CUDA >= 10.2), so every entry point is self-sufficient -- previously only finetune_ddp set it. DDP helpers: pretrain_ddp.py now imports configure_nccl_for_topology and ddp_setup from kermt/util/ddp_utils.py (the shared module this branch added for finetune) instead of keeping its own copies. All moved code is behavior-preserving (the duplicates were functionally identical). Also drops now-unused random/numpy imports. Signed-off-by: Eva Xue <evax@nvidia.com>
The per-batch loss was normalized by each rank's LOCAL valid-label count before backward, while DDP averages the per-rank gradients by 1/world_size. When per-rank counts differ (the sparse multi-task case) this weights ranks equally rather than per-label: a task whose labels land on only k of N ranks gets its gradient scaled by ~k/N vs a single-GPU run, so rare tasks are systematically down-weighted, and MTLLoss's log_sigma weights are biased the same way. DDP finetuning therefore did not match single-GPU. Fix both the MTL and pooled branches: all-reduce the valid-label COUNT (denominator) across ranks while keeping the loss SUM (numerator) local, then scale by world_size so DDP's 1/world_size average lands on the true global per-label mean. This matches single-GPU exactly and reduces to the previous behavior when world_size == 1 or per-rank counts are equal. The existing manual log_sigma gradient averaging then also yields the correct single-GPU gradient. Signed-off-by: Eva Xue <evax@nvidia.com>
drop_last: the DDP train DataLoader used drop_last=True, discarding up to world_size*(batch_size-1) samples -- a large fraction for small finetune sets on many GPUs -- and diverging from the single-GPU loader (which keeps the last batch). The DistributedSampler already pads to an equal per-rank sample count, so drop_last=False is DDP-safe (equal batch counts keep ranks in sync) and uses all the data. validation: every rank was scoring the full val set redundantly (N x compute, which hurts for large val sets under DDP). Now only rank 0 evaluates the unwrapped model (a plain non-DDP forward -- no collectives, safe on one rank) and broadcasts val_scores/val_loss so all ranks make identical best-model / early-stop decisions and stay in sync. Signed-off-by: Eva Xue <evax@nvidia.com>
The _about_gpu_selection note suggested "--gpus 0,2 etc." for multiple GPUs, but run_finetune_local.py resolves --gpus to a single device (multi-id is rejected) since multi-GPU moved to --num-gpus (DDP). Clarify: --gpus <id> picks a single device; --num-gpus N is data-parallel. Signed-off-by: Eva Xue <evax@nvidia.com>
Remove finetune_ddp.py and fold data-parallel finetuning into the single main.py finetune entrypoint. DDP is selected at runtime by the WORLD_SIZE environment variable: unset -> single-process path (unchanged); set -> spawn one process per GPU via mp.spawn, each running cross_validate with the DDP rank/world_size. The parsed args Namespace is passed straight to the worker (picklable), avoiding the old launcher's sys.argv re-parse. WORLD_SIZE is validated against the visible GPU count. WORLD_SIZE=1 exercises the DDP path on a single GPU. --batch_size stays per-GPU under DDP, matching pretrain_ddp.py. The finetune skill runner (run_finetune_local.py) now always builds a main.py finetune argv and sets WORLD_SIZE=num_gpus for multi-GPU; SKILL.md and the shared determinism docstring are updated to match. No mirror finetune-side launcher is kept. Signed-off-by: Eva Xue <evax@nvidia.com>
758034b to
675d02e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds opt-in data-parallel (DistributedDataParallel) finetuning so finetunes can be accelerated across multiple GPUs, mirroring the existing pretraining DDP path. Single-GPU finetuning is unchanged and remains the default — DDP only
engages through a dedicated launcher.
What's new
Core (
feat(finetune): add optional multi-GPU DDP training)finetune_ddp.py: a one-process-per-GPU launcher (torch.multiprocessing spawn,WORLD_SIZEenv, NCCL topology auto-config, per-rank determinism), routed through the normalcross_validate -> run_trainingfinetune path.kermt/util/ddp_utils.py: sharedddp_setup+configure_nccl_for_topologyso the pretrain and finetune launchers use one DDP bootstrap.task/train.py(run_training/train): threadrank/world_size; wrap the model in DDP (find_unused_parameters=True, since the finetune model carries parameters the FFN forward does not exercise); shard the train loader with aDistributedSampler(+set_epoch); run validation on the unwrapped module over the full val set on every rank so best-model selection stays consistent; gate checkpoint saving, logging and test-set evaluation to rank 0; all-reduce the standaloneMTLLoss.log_sigmagradient. Also uses a distincttrain_loaderinstead of reassigningtrain_data, fixing anensemble_size>1DataLoader double-wrap.cross_validate: threadrank/world_size; gate W&B + reporting to rank 0.build_lr_scheduler: world-size-awaresteps_per_epoch.Skill (
feat(kermt-finetune skill): expose optional multi-GPU DDP finetune)agent/scripts/run_finetune_local.py: new--num-gpus Noption. Default 1 (single-processmain.py finetune, unchanged);N>1launchesfinetune_ddp.pywithWORLD_SIZE=N. Under DDP it omits the single-device--gpuflag and recordsnum_gpus/distributedinrun.json.agent/skills/kermt-finetune/SKILL.md: documents--num-gpusin the hardware, arguments, invocation-template, and error-hint sections.Design
finetune_ddp.py).python main.py finetune ...and--num-gpus 1behave exactly as before.--batch_sizeis per process; effective global batch isbatch_size * WORLD_SIZE. The LR schedule is world-size-aware, so the same--batch_sizecan be reused as GPUs are added.Usage
Testing
main.py finetune) runs unchanged end-to-end (2-epoch regression finetune from a pretrained checkpoint): trains, selects best by validation, saves checkpoints, evaluates on test.finetune_ddp.py,WORLD_SIZE=1): full DDP wrap, DistributedSampler +set_epoch, unwrapped-module eval, rank-0 checkpoint/test — completes with comparable validation/test metrics to the single-process run.--num-gpus 1dispatches tomain.py finetune;--num-gpus N(N>1) dispatches tofinetune_ddp.pywithWORLD_SIZE=N(verified by dry-run command inspection and an actual launch). Existing finetune-runner unit tests (19) pass.