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
6 changes: 4 additions & 2 deletions megatron/model/fused_layer_norm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,8 +87,10 @@ def reset_parameters(self):
def forward(self, input):
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
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
Copy link
Member Author

@thomasw21 thomasw21 Mar 25, 2022 •
edited
Loading

Choose a reason for hiding this comment

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

@stas00

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.

Copy link
Contributor

Choose a reason for hiding this comment

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

Should we extend:

Megatron-DeepSpeed/megatron/mpu/mappings.py

Lines 22 to 30 in 87a9dba

def _reduce(input_):
"""All-reduce the the input tensor across model parallel group."""
# Bypass the function if we are using only 1 GPU.
if get_tensor_model_parallel_world_size()==1:
return input_
# All-reduce.
torch.distributed.all_reduce(input_, group=get_tensor_model_parallel_group())

to support an optional ReduceOp.AVG

Copy link
Member Author

Choose a reason for hiding this comment

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

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).

Copy link
Member Author

Choose a reason for hiding this comment

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

Also I don't think we save much compute by supporting that.


return FusedLayerNormAffineFunction.apply(
input, weight, bias, self.normalized_shape,self.eps)
Expand Down