Skip to content

Add Zero-offload support #56

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
siddharth9820 merged 1 commit into from
Jul 7, 2022
Merged

Add Zero-offload support #56

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 24 additions & 15 deletions megatron/optimizer/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,27 +49,36 @@ def _get_params_for_weight_decay_optimization(modules):
def get_megatron_optimizer(model):
args = get_args()

if args.cpu_optimizer:
raise NotImplementedError('need to add cpu adam')

# Base optimizer.
param_groups = _get_params_for_weight_decay_optimization(model)
if args.create_moe_param_group:
from deepspeed.moe.utils import is_moe_param, split_params_into_different_moe_groups_for_optimizer
param_groups = split_params_into_different_moe_groups_for_optimizer(param_groups)
if args.optimizer == 'adam':
optimizer = Adam(param_groups,
lr=args.lr,
weight_decay=args.weight_decay,
betas=(args.adam_beta1, args.adam_beta2),
eps=args.adam_eps)
elif args.optimizer == 'sgd':
optimizer = SGD(param_groups,
lr=args.lr,
weight_decay=args.weight_decay,
momentum=args.sgd_momentum)

if args.cpu_optimizer:
assert args.optimizer == 'adam', 'CPU offloading is for Adam'
if args.cpu_torch_adam:
cpu_adam_optimizer = torch.optim.AdamW
else:
from deepspeed.ops.adam import DeepSpeedCPUAdam
cpu_adam_optimizer = DeepSpeedCPUAdam
optimizer = cpu_adam_optimizer(param_groups,
lr=args.lr,
weight_decay=args.weight_decay)
else:
raise Exception('{} optimizer is not supported.'.format(
if args.optimizer == 'adam':
optimizer = Adam(param_groups,
lr=args.lr,
weight_decay=args.weight_decay,
betas=(args.adam_beta1, args.adam_beta2),
eps=args.adam_eps)
elif args.optimizer == 'sgd':
optimizer = SGD(param_groups,
lr=args.lr,
weight_decay=args.weight_decay,
momentum=args.sgd_momentum)
else:
raise Exception('{} optimizer is not supported.'.format(
args.optimizer))

if args.deepspeed:
Expand Down