Skip to content

fix(trainer): correct LR schedule, final checkpoint and STG under DDP - #263

Open
blue (Blueyyds) wants to merge 1 commit into
Lightricks:mainfrom
wgdl666:fix/trainer-multi-gpu-correctness
Open

fix(trainer): correct LR schedule, final checkpoint and STG under DDP#263
blue (Blueyyds) wants to merge 1 commit into
Lightricks:mainfrom
wgdl666:fix/trainer-multi-gpu-correctness

Conversation

@Blueyyds

Copy link
Copy Markdown

Three defects that only surface on multi-GPU runs or with checkpoint retention enabled:

  1. LR scheduler advanced once per process. Accelerate's default step_scheduler_with_optimizer=True applies its distributed batch-size adjustment on top of the trainer's own explicit scheduler.step() per optimizer step, so an 8-GPU run consumed the schedule 8x too fast and hit the final LR after 1/8 of training. The trainer owns the stepping, so opt out.

  2. Final checkpoint could delete itself. When the last optimizer step also lands on a save interval, the same path was written twice and appended to the retention list twice; with keep_last_n=1 the pruning pass then removed the file it had just written. Skip the redundant save when the interval checkpoint already exists, with a wait_for_everyone() so all ranks agree on the state.

  3. STG validation crashed under DDP. transformer.num_blocks is not reachable through the DistributedDataParallel wrapper, so enabling stg_scale raised AttributeError during validation. Read the attribute off .module when wrapped.

Three defects that only surface on multi-GPU runs or with checkpoint
retention enabled:

1. LR scheduler advanced once per process. Accelerate's default
   step_scheduler_with_optimizer=True applies its distributed
   batch-size adjustment on top of the trainer's own explicit
   scheduler.step() per optimizer step, so an 8-GPU run consumed the
   schedule 8x too fast and hit the final LR after 1/8 of training.
   The trainer owns the stepping, so opt out.

2. Final checkpoint could delete itself. When the last optimizer step
   also lands on a save interval, the same path was written twice and
   appended to the retention list twice; with keep_last_n=1 the
   pruning pass then removed the file it had just written. Skip the
   redundant save when the interval checkpoint already exists, with a
   wait_for_everyone() so all ranks agree on the state.

3. STG validation crashed under DDP. transformer.num_blocks is not
   reachable through the DistributedDataParallel wrapper, so enabling
   stg_scale raised AttributeError during validation. Read the
   attribute off .module when wrapped.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant