Skip to content

Commit

Permalink
Fix #726 grad checkpointing + dynamo work in newer torch, but w/ ddp,…
Browse files Browse the repository at this point in the history
… ddp optimizer must be disabled
  • Loading branch information
rwightman committed Oct 30, 2024
1 parent 82d7496 commit 49eac2f
Showing 1 changed file with 6 additions and 0 deletions.
6 changes: 6 additions & 0 deletions src/open_clip_train/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -422,6 +422,12 @@ def main(args):
original_model = model
if args.torchcompile:
logging.info('Compiling model...')

if args.grad_checkpointing and args.distributed:
logging.info('Disabling DDP dynamo optimizer when grad checkpointing enabled.')
# As of now (~PyTorch 2.4/2.5), compile + checkpointing but DDP optimizer must be disabled
torch._dynamo.config.optimize_ddp = False

model = torch.compile(original_model)

if 'train' not in data:
Expand Down

0 comments on commit 49eac2f

Please sign in to comment.