ZeRO & Custom Optmizer (RangerLars)

[CalogeroZarbo]

Hi,
first of all thanks for the contribution. I'm trying to use DeepSpeed and Zero with a custom optimizer from https://github.com/mgrankin/over9000
More in the specific I'm using the implementation of the RangerLars (RAdam+LARS+LookAhead), which in that repo has been called Over9000.
Is there any way to wrap a custom optimizer like this in a Zero compatible one? I found this class in your repository:

class FP16_DeepSpeedZeroOptimizer(object):
    """
    DeepSpeedZeroOptimizer designed to reduce the memory footprint
    required for training large deep learning models.
    For more details please see ZeRO: Memory Optimization Towards Training A Trillion Parameter Models
    https://arxiv.org/abs/1910.02054
    For usage examples, refer to TODO: DeepSpeed V2 Tutorial
    """

Can I use it for my purpose? If yes, is it possible to have a snippet of code?

Thank you,
Cal

[tjruwase]

Thanks for your interest in DeepSpeed. In theory it should be possible to wrap RangerLars with ZeRO. But there are a couple of things to consider:

1. RangerLars itself seems to be a wrapper like this: LookAhead(RAdam). And it seems you want to do an extra level of wrapping like this: ZeRO(LookAhead(RAdam))). We have not previously explored this multi-level wrapping.
2. However, a quick skim of ZeRO and LookAhead apis suggests that ZeRO(LookAhead(...)) is possible .

So a quick way to try this out is the following three steps:

1. Update your json config to enable fp16and ZeRO

2. Make sure there is no optimizer field in your json config.

3. Initialize RangerLars as desired in your client script and pass to deepspeed for ZeRO wrapping through the deepspeed.initialize() call.

Please share your experience, I am especially curious :).

[CalogeroZarbo]

Hi @tjruwase, thank you for your quick response.

This is the stack trace:

  File "train_seq2seq.py", line 447, in main
    pp_rank = torch.distributed.get_rank(group=self.optimizer.dp_process_group)
AttributeError: 'FP16_UnfusedOptimizer' object has no attribute 'dp_process_group'
    zero_optimization=zero_optimization
  File "train_seq2seq.py", line 80, in train_encdec_v1
    encoder_engine, encoder_optimizer, trainloader, _ = deepspeed.initialize(args=cmd_args, model=encoder, optimizer=encoder_optimizer, model_parameters=encoder_params, training_data=train_dataset, dist_init_required=True)
  File "/home/calogero_zarbo/miniconda3/envs/virus_torch/lib/python3.6/site-packages/deepspeed/__init__.py", line 96, in initialize
    collate_fn=collate_fn)
  File "/home/calogero_zarbo/miniconda3/envs/virus_torch/lib/python3.6/site-packages/deepspeed/pt/deepspeed_light.py", line 182, in __init__
    self._configure_checkpointing(dist_init_required)
  File "/home/calogero_zarbo/miniconda3/envs/virus_torch/lib/python3.6/site-packages/deepspeed/pt/deepspeed_light.py", line 337, in _configure_checkpointing
    pp_rank = torch.distributed.get_rank(group=self.optimizer.dp_process_group)
AttributeError: 'FP16_UnfusedOptimizer' object has no attribute 'dp_process_group'

I looked for the dp_process_group function and I thought I need to wrap all of this somehow. Maybe it could work to have: ZeRO(LookAhead(ZeRO(RAdam))?

[tjruwase]

Thanks for exploring this issue. The failure you are observing is happening during DeepSpeed initialization. The problem in this case is that we have the following check in deepspeed_light to restrict ZeRO to wrapping only Adam. But for your case, can you please disable the Adam restriction so that ZeRO can wrap any optimizer, and then we can see what happens?

https://github.com/microsoft/DeepSpeed/blob/1496247a19d4bfd8346001ac491b1d160b5ecbee/deepspeed/pt/deepspeed_light.py#L444

[CalogeroZarbo]

Hi @tjruwase!
I modified the link of code as you suggested in this way:

if self.zero_optimization(): # and self.optimizer_name() == ADAM_OPTIMIZER: 

And it worked! Now I'm training the model using ZeRO enabled.

Thank you!

PS. In order to make this available for all should we plan a PR or it will be enabled by you in the future?

[tjruwase]

@jeffra, @samyam, @ShadenSmith

Yay!!! This is excellent news!!!

Did you verify correctness for your model? Correctness is the main concern that we have for restricting to Adam which we have tested internally. But perhaps we just need to generate a warning for brave users such as yourself to alert you that you are truly pushing the frontiers :)

You can prepare the PR. In the PR please add a warning such as below here,

if self.zero_optimization() and self.optimizer_name != ADAM_OPTIMIZER:
   logging.warning("**** You are using ZeRO with an untested optimizer, proceed with caution *****")

Thanks for doing this great service.

[CalogeroZarbo]

Good morning @tjruwase!

I did some tests and it seems that the specific implementation of RangerLars that I'm using needs to be further investigated to make it work properly as without ZeRO enabled.

The training starts, the loss is backpropagated but the model is not learning (the loss stays around the same value of the first step).
My understanding is that: technically we can use custom optimizers, but we need to be sure ourselves that the code is suitable for ZeRO.

I would do the PR anyway to be all of us able to test our models with different optimizers and, as a result of our tests, create a little guideline on how to create custom optimizers compliant with ZeRO.

What would you think?

Thanks for doing this great service.

Thank you for your support, It is motivating being part of this community and contributing to such an interesting project.

Cheers,
Cal