Sync layer norm #271
Sync layer norm #271
Conversation
thomasw21
changed the base branch from
main
to
thomas/test_different_layer_norm
looks like pt-1.11 dropped `torch.testing.assert_equal`, so using `torch.testing.assert_equal` instead
tests/test_training.py
Outdated
| # 2. test training from checkpoint: resume | ||
| # now do it again, this time resuming from the checkpoint | ||
| with CaptureStdout() as cs: | ||
| execute_subprocess_async(cmd, env=self.get_env()) |
There was a problem hiding this comment.
so it crashes on resume:
Traceback (most recent call last):
File "/home/stas/anaconda3/envs/py38-pt111/lib/python3.8/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 345, in wrapper
return f(*args, **kwargs)
File "/mnt/nvme0/code/huggingface/Megatron-DeepSpeed-master-4/pretrain_gpt.py", line 245, in main
pretrain(train_valid_test_datasets_provider, model_provider, forward_step,
File "/mnt/nvme0/code/huggingface/Megatron-DeepSpeed-master-4/megatron/training.py", line 188, in pretrain
iteration = train(forward_step_func,
File "/mnt/nvme0/code/huggingface/Megatron-DeepSpeed-master-4/megatron/training.py", line 857, in train
train_step(forward_step_func,
File "/mnt/nvme0/code/huggingface/Megatron-DeepSpeed-master-4/megatron/training.py", line 441, in train_step
loss = model[0].train_batch(data_iter=data_iterator)
File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/pipe/engine.py", line 346, in train_batch
self._exec_schedule(sched)
File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/pipe/engine.py", line 1363, in _exec_schedule
self._exec_instr(**cmd.kwargs)
File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/pipe/engine.py", line 1149, in _exec_optimizer_step
self._take_model_step(lr_kwargs)
File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/engine.py", line 1787, in _take_model_step
self.optimizer.step()
File "/home/stas/anaconda3/envs/py38-pt111/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/bf16_optimizer.py", line 239, in step
assert all_groups_norm > 0.
AssertionError
megatron/model/fused_layer_norm.py
Outdated
| tp_world_size = mpu.get_tensor_model_parallel_world_size() | ||
| # TODO: hack in order to synchronize all layer norms despite them being unsynched | ||
| weight = mpu.reduce_from_tensor_model_parallel_region(self.weight) / tp_world_size | ||
| bias = mpu.reduce_from_tensor_model_parallel_region(self.bias) / tp_world_size | ||
|
|
||
| return FusedLayerNormAffineFunction.apply( | ||
| input, self.weight, self.bias, self.normalized_shape,self.eps) | ||
| input, weight, bias, self.normalized_shape,self.eps) |
There was a problem hiding this comment.
@tjruwase, this is the main workaround that does the all_reduce mean on layer norm's weight+bias that we want to put in until we can fix the fp32 weights.
megatron/model/fused_layer_norm.py
Outdated
| weight = torch.clone(self.weight) | ||
| bias = torch.clone(self.bias) | ||
| weight = mpu.reduce_from_tensor_model_parallel_region(weight) / tp_world_size | ||
| bias = mpu.reduce_from_tensor_model_parallel_region(bias) / tp_world_size |
There was a problem hiding this comment.
Essentially the reduce is an in-place operator, which means at each forward pass, self.weight was updated with the sum of all the weights of all tp_ranks. We could try thinking of a better fix by doing a average reduce, but I'm scared back propagation doesn't play well with this in place logic.
New test fails with:
E raise StopIteration
E StopIteration
This is more expected since the previous run should have consumed all the tokens. Going to update #272 and restart the training.
There was a problem hiding this comment.
Should we extend:
Megatron-DeepSpeed/megatron/mpu/mappings.py
Lines 22 to 30 in 87a9dba
to support an optional ReduceOp.AVG
There was a problem hiding this comment.
I think this is tricky. The reason why is this means that we need to implement custom backward function (since you compute the average, the gradient needs to be divided by the tp world size).
There was a problem hiding this comment.
Also I don't think we save much compute by supporting that.
* Enable universal ckpting * Update run scripts * Address PR feedback * Remove line * Fix white lines * Remove redudant changes * Apply to gpt_model only * Code cleanup * Code cleanup * Update training.py Co-authored-by: Michael Wyatt <mrwyattii@gmail.com> * Update training.py Co-authored-by: Michael Wyatt <mrwyattii@gmail.com> * Log loss_scale only valid for fp16 * Add README and bf16 scripts * Visualization docsts * Support older DS * Handle uni_ckpt import error * Revert changes --------- Co-authored-by: Michael Wyatt <mrwyattii@gmail.com>
Force sync layer norms