Repository: hpcaitech/ColossalAI

Issue:
[BUG]: HybridParallelOptimizer holds unsharded model parameters after sharding
### 🐛 Describe the bug

When using tensor parallelism, model parameters are sharded across GPUs to reduce its memory consumption and parallel execution.
However, the optimizer still holds unsharded model parameters, preventing the old unsharded parameters from being released, taking more memory.

Example code: (adopted from `examples/language/gpt2/hybridparallelism/finetune.py`)
```python
colossalai.launch_from_torch(config={})
plugin = HybridParallelPlugin(tp_size=4, pp_size=1)
optimizer = Adam(model.parameters())
# initialize dataloader

model, optimizer, *_ = booster.booster(model, optimizer, ...)
```
```
> model.module.transformer.wte.weight
Parameter containing:
tensor([[-0.1101, -0.0393, ...]], device='cuda:0', dtype=torch.float16, requires_grad=True)

> model.module.transformer.wte.weight.shape
torch.Size([12565, 768])

> optimizer.param_groups[0]["params"][0]
Parameter containing:
tensor([[-0.1101, -0.0393, ...]], device='cuda:0', requires_grad=True)

> optimizer.param_groups[0]["params"[0].shape
torch.Size([50257, 768])
```

This also affects `MixedPrecisionOptimizer.master_to_working_map` and `MixedPrecisionOptimizer.working_to_master_map`:
```
# model.module.transformer.wte.weight is supposed to be in a working parameter
> model.module.transformer.wte.weight.shape
torch.Size([12565, 768])
> id(model.module.transformer.wte.weight)
139684649437120

# First working parameter in map does not refer to this
> list(iter(optimizer.master_to_working_map))[0].shape
torch.Size([50257, 768])
> id(list(iter(optimizer.master_to_working_map))[0])
139693862695728
```

Because of this it seems only a portion of parameters (ie. unsharded ones) only trained, as `MixedPrecisionOptimizer.step()` skips sharded parameters as gradients are not stored in mismatched unsharded parameters:

https://github.com/hpcaitech/ColossalAI/blob/df5e9c53cf23d44656470cc319ee0b470c40712f/colossalai/amp/naive_amp/mixed_precision_optimizer.py#L173-L175

### Environment

PyTorch 2.2.1 / CUDA 12.1


Discussion:
Hi, thanks for the issue. 
I reproduced the bug using this script
[finetune.zip](https://github.com/hpcaitech/ColossalAI/files/14821735/finetune.zip)
This might be due to some unexpected model movement without ZeRO. Mostly ZeRO is used and the params are sharded in-place. I'm looking into this.

Resolving patch:
diff --git a/colossalai/shardformer/shard/shard_config.py b/colossalai/shardformer/shard/shard_config.py
--- a/colossalai/shardformer/shard/shard_config.py
+++ b/colossalai/shardformer/shard/shard_config.py
@@ -74,8 +74,10 @@ def _turn_on_all_optimization(self):
         self.enable_fused_normalization = True
         self.enable_flash_attention = True
         self.enable_jit_fused = True
-        self.enable_sequence_parallelism = True
-        self.enable_sequence_overlap = True
+        # This can cause non-in-place param sharding when used without ZeRO.
+        # It may also slow down training when seq len is small. Plz enable manually.
+        # self.enable_sequence_parallelism = True
+        # self.enable_sequence_overlap = True
 
     def _infer(self):
         """


Test patch: