Fix ZeRO-3 generation context manager

[lewtun]

This PR incorporates a fix from @pacman100 to enable the unwrap_model_for_generation() context manager to work with other trainers like DPO.

The basic issue with our original implementation is that we weren't removing all active parameters during the hook removal and this led to DeepSpeed errors like:

AssertionError: {'id': 0, 'status': 'AVAILABLE', 'numel': 25755648, 'ds_numel': 25755648, 'shape': (50304, 512), 'ds_shape': (50304, 512), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': {182}, 'ds_tensor.shape': torch.Size([3219456])}

Tested with this gist: https://gist.github.com/pacman100/a6c89a681f8f76bdf17f5bf6874eb983

cc @edbeeching you should rebase on main once this is merged so #1605 will work

Closes #1543

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[edbeeching]

Thanks, LGTM

[younesbelkada]

Thanks a lot !