Skip to content

GitLab

Unverified Commit 7a065dc9 authored by HELSON's avatar HELSON Committed by GitHub
Browse files

[hotfix] fix megatron_init in test_gpt2.py (#1357)

parent 644582ee
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
tests/test_tensor/model/test_gpt2.py
View file @ 7a065dc9
...@@ -56,7 +56,7 @@ def init_megatron_spec(model, pg: ProcessGroup): ...@@ -56,7 +56,7 @@ def init_megatron_spec(model, pg: ProcessGroup):
elif 'wte' in mn or 'wpe' in mn: elif 'wte' in mn or 'wpe' in mn:
assert 'weight' in pn assert 'weight' in pn
split_param_col_tp1d(param, pg) split_param_col_tp1d(param, pg)
elif 'c_fc' in mn or 'c_proj' in mn: elif 'c_attn' in mn or 'c_proj' in mn:
split_param_col_tp1d(param, pg) split_param_col_tp1d(param, pg)
# debug_print([0], '\t', param.compute_spec, param.shape) # debug_print([0], '\t', param.compute_spec, param.shape)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment