Add DeepSpeed support

[thevasudevgupta]

This PR will add DeepSpeed support to Accelerate.

User's code should look like this:

from accelerate import Accelerator

model: torch.nn.Module
optimizer: torch.optim.Optimizer
tr_data: torch.utils.data.DataLoader
eval_data: torch.utils.data.DataLoader

gradient_accumulation_steps: int
epochs: int

accelerator = Accelerator(fp16=True)
model, optimizer, tr_data, eval_data = accelerator.prepare(model, optimizer, tr_data, eval_data)

# training loop
for e in range(epochs):
    for step, batch in enumerate(tr_data):
        output = model(**batch)
        loss = output.loss / gradient_accumulation_steps
        accelerator.backward(loss)
        if (step + 1) % gradient_accumulation_steps == 0 or step == len(train_data) - 1:
            optimizer.step()
            optimizer.zero_grad()
-           lr_scheduler.step()
+           if not optimizer.is_overflow:
+               lr_scheduler.step()

# distributed evaluation
model.eval()
for batch in eval_data:
    output = model(**batch)
    predictions = accelerator.gather(output)
    labels = accelerator.gather(batch["labels"])

# for saving your model
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)

# if its a HuggingFace Transformer
- unwrapped_model.save_pretrained(args.output_dir, save_function=accelerator.save)
+ unwrapped_model.save_pretrained(args.output_dir, state_dict=accelerator.get_state_dict(model), save_function=accelerator.save)

Running the script

accelerate config
accelerate launch <training_script.py> <training_args>

@sgugger

[sgugger]

Thanks a lot for your work on this! This is great! I left a couple of comments to polish the implementation.

[sgugger]

Note for before merging this: the PR should be rebased on master and the reference to the model added in the Accelerator should be cleaned up in the method introduced in the PR above.

[sgugger]

Thanks a lot for all your amazing work! I think we are in good shape to have this merge, I just left some last comments around naming.

[LysandreJik]

Cool, great job implementing that @vasudevgupta7!

Two items that should be handled in my opinion:

• It seems that accelerate will not run once this PR is merged if deepspeed isn't in the environment. Is this correct?
• There really should be some tests to ensure that everything is working correctly. I would advocate for a test suite that runs accelerate with and without deepspeed, but I know the infra to support that isn't there yet.

[sgugger]

Just a few last tweaks to the doc and this should be good to merge!