[zero] add zero optimizer for ColoTensor #1046
Conversation
| UNSCALED = 1 | ||
|
|
||
|
|
||
| class ZeroOptimizer(ColossalaiOptimizer): |
There was a problem hiding this comment.
Should this be put in the zero module?
tests/test_tensor/test_zero_optim.py
Outdated
| torch_model.train() | ||
| set_seed(gpc.get_local_rank(ParallelMode.DATA)) | ||
| for i, (input_ids, attn_mask) in enumerate(train_dataloader): | ||
| if i > 1: |
There was a problem hiding this comment.
why only one pass? will it fail if i > 5?
There was a problem hiding this comment.
This test takes too long if i > 5
There was a problem hiding this comment.
OK, but you should have more robust test like i > 3
colossalai/zero/zero_optimizer.py
Outdated
| max_scale=max_scale) | ||
| self._found_overflow: torch.Tensor = torch.zeros(1, dtype=torch.int64, device=torch.cuda.current_device()) | ||
| self.dp_process_group = gpc.get_group(ParallelMode.DATA) | ||
| self.mp_process_group = gpc.get_group(ParallelMode.MODEL) |
There was a problem hiding this comment.
why don't we have a ParallelMode.Global?
There was a problem hiding this comment.
@FrankLeeeee shall we add this mode?
There was a problem hiding this comment.
I think implementing it is not difficult. You can consider it as a patch in future.
There was a problem hiding this comment.
We have this mode actually...
| for p in group['params']: | ||
| if not self.module.chunk_manager.is_chunk_free(p): | ||
| fp32_p = self.fp16_param_to_fp32_param[p] | ||
| self.module.chunk_manager.copy_tensor_to_chunk_slice(p, fp32_p) |
There was a problem hiding this comment.
copy in tensors is inefficient....
You should add a TODO and improve this line later?
tests/test_tensor/test_gpt.py
Outdated
| @@ -68,6 +68,7 @@ def run_gpt(init_spec_func, use_ddp): | |||
| for i, (input_ids, attn_mask) in enumerate(train_dataloader): | |||
| logits = model(input_ids, attn_mask) | |||
| torch_logits = torch_model(input_ids, attn_mask) | |||
| print(torch_logits, logits) | |||
No description provided.