Apex optimizer interface differs from torch.nn.Optim

Open
#1,014 2 comments 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
25/100
Issue type
Refactor
Clarity
Needs clarification
Activity status
Stale
Tech stack
python, pytorch
Domain
machine-learning

Research direction

Start by comparing the Apex/core optimizer interface with the OSS optimizer interface, focusing on zero_grad(), initialization, checkpoint loading, and optimizer state representation. Done means the interfaces are consistently usable by generic training frameworks, including gradient clearing and checkpoint save/load behavior.

Written by the indexing model from the issue text.

Description

This makes it harder to interchangeably use optimizers in generic training frameworks (e.g. checkpointing/loading from checkpoint has to be different depending on the optimizer, zeroing/setting grads to None is different etc). Some examples:

  • oss optim's set_grad_none is in the zero_grad(), and apex's set_grad_none is in the init().
  • apex's optimizer state are lazily initiated which makes checkpoint loading a bit tricky
  • apex and oss optim has mismatch in the state, e.g. "step" can either be a tensor or int?

Is it possible to make apex/core optimizer interface more consistent?
cc @ptrblck
cc @xw285cornell

Dominant language
Python
Stars
9k
Forks
1.5k
Avg merge
2d 4h
Merged PRs (30d)
3

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from NVIDIA/apex

All issues in NVIDIA/apex

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.