[draft] proposed fix for incorrect mask application in FSDP

[bfineran]

in the current layer masking implementation, parameters are stored at initialization and references to them are used to apply masks on modifier update

in FSDP mode, it seems that masking on top of these references does update these parameter references, but these parameters no longer have an effect on the FSDP module.

this fix just implements the simple flow @dsikka used in the sparsify MVP to apply the masks over a reference to the current FSDP module at update time. (ie instead of applying masks on the saved references to the layers, the masks are applied directly over fresh references from the model)

handing off to @Satrat

confirmation of fix
test command:
accelerate launch --config_file fsdp_config.yaml test_trainer.py

snippet of output with sparsity log (previously 0.0 for all sparsity values):

2023-10-31 16:22:18 sparseml.transformers.finetune.session_mixin INFO     Finalized SparseML recipe argument applied to the model
2023-10-31 16:22:18 sparseml.transformers.finetune.session_mixin INFO     Sparsification info for ./obcq_deployment: 15191712 total params. Of those there are 15187968 prunable params which have 15.291973225121358 avg sparsity.
2023-10-31 16:22:18 sparseml.transformers.finetune.session_mixin INFO     sparse model detected, all sparsification info: {"params_summary": {"total": 15191712, "sparse": 2322540, "sparsity_percent": 15.288204515725418, "prunable": 15187968, "prunable_sparse": 2322540, "prunable_sparsity_percent": 15.291973225121358, "quantizable": 15187968, "quantized": 0, "quantized_percent": 0.0}, "params_info": {"_fsdp_wrapped_module.model.layers.0.self_attn.q_proj.weight": {"numel": 82944, "sparsity": 0.5000361800193787, "quantized": false}, "_fsdp_wrapped_module.model.layers.0.self_attn.k_proj.weight": {"numel": 82944, "sparsity": 0.5000361800193787, "quantized": false}, "_fsdp_wrapped_module.model.layers.0.self_attn.v_proj.weight": {"numel": 82944, "sparsity": 0.5000361800193787, "quantized": false}, "_fsdp_wrapped_module.model.layers.0.self_attn.o_proj.weight": {"numel": 82944, "sparsity": 0.5000361800193787, "quantized": false}, "_fsdp_wrapped_module.model.layers.0.mlp.gate_proj.weight": {"numel": 221184, "sparsity": 0.5000135898590088, "quantized": false}, "_fsdp_wrapped_module.model.layers.0.mlp.up_proj.weight": {"numel": 221184, "sparsity": 0.5000135898590088, "quantized": false}, "_fsdp_wrapped_module.model.layers.0.mlp.down_proj.weight": {"numel": 221184, "sparsity": 0.0, "quantized": false}, "_fsdp_wrapped_module.model.layers.1.self_attn.q_proj.weight": {"numel": 82944, "sparsity": 0.5000361800193787, "quantized": false}, "_fsdp_wrapped_module.model.layers.1.self_attn.k_proj.weight": {"numel": 82944, "sparsity": 0.5000361800193787, "quantized": false}, "_fsdp_wrapped_module.model.layers.1.self_attn.v_proj.weight": {"numel": 82944, "sparsity": 0.5000361800193787, "quantized": false}, "_fsdp_wrapped_module.model.layers.1.self_attn.o_proj.weight": {"numel": 82944, "sparsity": 0.5000361800193787, "quantized": false}, "_fsdp_wrapped_module.model.layers.1.mlp.gate_proj.weight": {"numel": 221184, "sparsity": 0.5000135898590088, "quantized": false}, "_fsdp_wrapped_module.model.layers.1.mlp.up_proj.weight": {"numel": 221184, "sparsity": 0.5000135898590088, "quantized": false},

Update 11/1/23

The above fix works for n_gpu=1 but multi-gpu. The latest commit should fix the issue with multi-gpu. Essentially what was happening is we were initializing the model to our SparseSession before it was wrapped by FSDP. To fix this I added a new callback for on_train_begin that replaces the session's pytorch model with the FSDP wrapped one.

Using model.apply as implemented in the initial fix works because FSDP overrides the module apply function, see https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.apply.

In order to access and update the underlying model outside of apply, we need to use the summon_full_params function, see https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params. This fixed the issue with reading out the sparsities after training

with FullyShardedDataParallel.summon_full_params(self.model):
    self.log_model_sparsification()

We may need to implement this idea in other areas of the codebase.

Remaining things to wrap up:

• currently this branch will only work with FSDP due to the summon_full_params call, update to check if we need this context before calling
• Having some issues with the script hanging when we try to save the model. Haven't had time to debug this, but https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading may be a good reference here. Might need to wrap this function so only one process is calling it?