[docs] critical API documentation is missing

[stas00]

It looks like things kind of work, except not quite right, and there are a lot of subtle nuances that are so hard to know about when integrating DeepSpeed. I think all these should be made loud and clear - and perhaps a simple full example of a training loop would help, including showing commented out code where the original training code is removed to do it the DeepSpeed-way.

As I am trying to figure out how to make gradient_accumulation_steps work correctly I'm finding all kinds of things I have missed when integrating DeepSpeed into HF Trainer. I will post them here as I find such things:

1. engine's backward returns loss, which it modifies under gradient_accumulation_steps > 1 but this is undocumented.

   • neither in the API docstring https://github.com/microsoft/DeepSpeed/blob/e60e92eb0a06673748c4cb63fbcf713ddd12fc22/deepspeed/runtime/engine.py#L852-L858
   • nor main docs: https://www.deepspeed.ai/getting-started/#training

2. Also it's not documented that the "client" shouldn't scale loss by gradient_accumulation_steps since Deepspeed does it in backward.

3. the fact that lr_scheduler.step happens inside engine's step is not documented in the API

   • https://github.com/microsoft/DeepSpeed/blob/e60e92eb0a06673748c4cb63fbcf713ddd12fc22/deepspeed/runtime/engine.py#L993-L996
   • but is documented in https://www.deepspeed.ai/getting-started/#training
   • it might be a good idea to also add an explicit - make sure to remove lr_scheduler.step() from your code if using DeepSpeed's scheduler.

4. the "client" must not skip engine.step() when gradient_accumulation_steps > 1, and since this is an integration of many methods this leads to a complicated brittle code:

                if self.deepspeed:
                    self.deepspeed.step()

                if (step + 1) % self.args.gradient_accumulation_steps == 0 or (
                    # last step in epoch but step is always smaller than gradient_accumulation_steps
                    steps_in_epoch <= self.args.gradient_accumulation_steps
                    and (step + 1) == steps_in_epoch
                ):
                    # Gradient clipping
                    if self.args.max_grad_norm is not None and self.args.max_grad_norm > 0 and not self.deepspeed:
                        # deepspeed does its own clipping
                        if self.use_amp:
                            # AMP: gradients need unscaling
                            self.scaler.unscale_(self.optimizer)
                        [...]
                        else:
                            # Revert to normal clipping otherwise, handling Apex or full precision
                            torch.nn.utils.clip_grad_norm_(
                                amp.master_params(self.optimizer) if self.use_apex else model.parameters(),
                                self.args.max_grad_norm,
                            )

                    # Optimizer step
                    if self.deepspeed:
                        pass # called outside the loop
                    [...]
                    else:
                        self.optimizer.step()

                    if not self.deepspeed:
                        self.lr_scheduler.step()

                    model.zero_grad()
                    [...]

After fixing the above 4 I managed to get the same weights and loss with bs=8/grad_accum=1 and bs=4/grad_accum=2. Yay!

[janEbert]

Adding to this:
5. Remove model = model.cuda() and model = model.half() calls.
6. Do keep data = data.cuda() and optionally add data = data.half() calls (see below comment for a dynamic option).
7. Substitute torch.nn.utils.clip_grad_norm_(params, X) with "gradient_clipping": X in the DeepSpeed config.
8. Guard parallel I/O using parallel primitives (except DeepSpeed checkpointing which is already nicely documented in the code and at deepspeed.ai).
9. Use torch.distributed methods for parallel primitives (e.g. torch.distributed.get_world_size(), torch.distributed.get_rank(), int(os.environ['LOCAL_RANK']), ...).

And a bit less interesting (though one of these should probably be adjusted instead of documented):
10. In deepspeed.initialize, the optimizer kwarg has precedence over the "optimizer" key in the DeepSpeed config.
11. In deepspeed.initialize, the "scheduler" key in the DeepSpeed config has precedence over the lr_scheduler kwarg.

[stas00]

Thank you for the additions, @janEbert!

6. Do keep data = data.cuda() and optionally add data = data.half() calls.

And since fp32 is being worked on #1004, this could be dynamic.

inputs = inputs.to(dtype=next(iter(model.parameters())).dtype)

so it can automatically handle the case of model.half() (fp16) or not (fp32) down the road.

In NLP models this usually happens automatically because the embedding converts the inputs to the right dtype.

[janEbert]

Oh nice, that input conversion is even better! Happy to hear FP32 ZeRO 3 is getting PRed as well. :)

[janEbert]

Actually I just noticed, so long as the device is set correctly (that is, with torch.cuda.set_device(local_rank)), model = model.cuda() can be left in and model = model.half() can be left in either way.

See device setup and model conversion code.

It would be nice to have a clear stance on whether these calls should be removed for future-proofing the code, though. In case there are changes to engine initialization.

[stas00]

Based on my understanding this is all taken care of internally by DeepSpeed, so you want to skip both of those calls if your training loop is running under Deepspeed - this is what we do in HF Trainer.

model = model.cuda() will OOM anyway if your model is larger than gpu0.

[janEbert]

Yea I guess it's better to just give your concrete advice; who knows how the internals will change with the speed this is moving at. :)