How to modify weights during training in a deepspeed stage 3 model

[lucasosouza]

This issue is more of a general question than a bug report: we are working with deepspeed stage 3 and model parallelization (training 10+ billion params LLMs). We need to be able to modify weights every N steps during training, but outside of the regular optimization process.

I've seen multiple ways of exporting and accessing weights from deepspeed models, but I haven't found any methods to modify them directly during training.

One option would be during saving/loading - but although I can load the weights using get_fp32_state_dict_from_zero_checkpoint, and possibly modify the state dict, I also can't find a way to reload this modified state dict into the model and continue training with the new set of weights.

I am currently using deepspeed as the strategy in Fabric (pytorch-lightning). Appreciate the help or any pointers.

[tjruwase]

@lucasosouza, thanks for your question. We have recently added supporting for examining various weight and optimizer state values, documented here. It sounds like you want the inverse functionality that can modify these values. Is that correct?

[tjruwase]

You might also find this documentation of BLOOM model training useful, in particular the following sections regarding on-the-fly fixes of model bugs:

1. https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md#2022-03-24-grad-clip-tp-sync-bug-fixing
2. https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md#2022-03-28-fixing-out-of-sync-layer-norms

[lucasosouza]

functionality

Yeah, exactly, we need to modify the weights. I am going through the sync_layer_norm example see if it applies to our use case. Would you know if there is any other available interface in deepspeed to modify the distributed weights (or instead modify a buffer registered to a nn.Module)?

[tjruwase]

There is no other interface in deepspeed for modifying distributed weights. The above is our first attempt to provide this kind of support. The current support does not even apply to registered buffers.

[tjruwase]

@lucasosouza, how are things going with this? Did you find the sync_layer_norm example useful? Thanks.

[tjruwase]

@lucasosouza, @dsj96, can you please provide feedback on the linked PR #4192? Thanks!

[lucasosouza]

PR #4192 solves for the use case we have! Thanks a lot

[lucasosouza]

Closing the issue given PR is merged