Skip to content

Sync layer norm #271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 38 commits into
base: thomas/test_different_layer_norm
Choose a base branch
from
Draft

Sync layer norm #271

Changes from 1 commit
Commits
Show all changes
38 commits
Select commit Hold shift + click to select a range
07ccb3d
Better
thomasw21 Mar 24, 2022
391ed48
Force synchronize the layer norms parameters across all TP
thomasw21 Mar 24, 2022
98d0e7c
import mpu
stas00 Mar 25, 2022
279a77e
use the bf16 branch for testing
stas00 Mar 25, 2022
87a9dba
`torch.testing.assert_equal` didn't make it (#273)
stas00 Mar 25, 2022
dbb5914
Merge remote-tracking branch 'origin/main' into thomas/fix_layer_norm
stas00 Mar 25, 2022
70f91f8
bf16 comms requite pt-1.11
stas00 Mar 25, 2022
835a3e5
already part of the function
stas00 Mar 25, 2022
37795a9
reproduce the crashing on resume
stas00 Mar 25, 2022
3ec65f7
run just the test we want for now
stas00 Mar 25, 2022
8271d41
all_reduce is an in_place operation
thomasw21 Mar 25, 2022
b418b47
Make a test that TP reshaping works
thomasw21 Mar 25, 2022
4b7207b
Woops
thomasw21 Mar 25, 2022
3bc5824
Woops
thomasw21 Mar 25, 2022
05c99db
Woops
thomasw21 Mar 25, 2022
55e10c6
Woops
thomasw21 Mar 25, 2022
2ab8a3a
Woops
thomasw21 Mar 25, 2022
d357839
Woops
thomasw21 Mar 25, 2022
5fb231c
Woops
thomasw21 Mar 25, 2022
cc7ff45
Woops
thomasw21 Mar 25, 2022
7cdb1be
Woops
thomasw21 Mar 25, 2022
4574ec9
Fix load issue
thomasw21 Mar 25, 2022
04e89d1
Woops
thomasw21 Mar 25, 2022
e943100
Fix checkpoint path
thomasw21 Mar 25, 2022
09cead3
Test that force sync will allow TP changes
thomasw21 Mar 25, 2022
77abee6
Nit
thomasw21 Mar 25, 2022
64a62c8
Now that we have a force sync mechanism, let's try to reproduce
thomasw21 Mar 29, 2022
0b7afcc
Compare model_states_rank
thomasw21 Mar 29, 2022
ce01733
test
thomasw21 Mar 29, 2022
89ab0b7
Row column bias should be synchronized as well
thomasw21 Mar 29, 2022
42997b2
New list of matching embeddings
thomasw21 Mar 29, 2022
e0ef168
Figure out why state differs
thomasw21 Mar 29, 2022
1fc4fe8
Test for final weight
thomasw21 Mar 29, 2022
7ebbed1
Test that torch_rng_state
thomasw21 Mar 29, 2022
2c49216
Fix non matching torch_rng_state for tp_rank=0
thomasw21 Mar 30, 2022
007ecb4
Update test
thomasw21 Mar 31, 2022
c3844b5
I'm surprised one can apply inplace operation here
thomasw21 Mar 31, 2022
189f054
Test out the loss from the fp32 weights and optimizer states
thomasw21 Apr 4, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions tests/test_training.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -706,3 +706,14 @@ def test_layer_norm_consistent(self, variation):
ref = weights[0]
for weight in weights[1:]:
torch_assert_equal(ref, weight, check_device=False)

# 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())
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


# test checkpoint loading
self.assertIn(f"successfully loaded checkpoint from {output_dir}/checkpoints", cs.out)

# test reports
self.assertIn("consumed samples", cs.out)