Skip to content

Commit

Permalink
[shardformer] update bloom/llama/vit/chatglm tests
Browse files Browse the repository at this point in the history 
[shardformer] update bloom/llama/vit/chatglm tests

[shardformer] update opt tests

[shardformer] update opt tests
  • Loading branch information
@flybird11111
flybird11111 committed Aug 13, 2023
1 parent 1e518ae commit 759646c
Show file tree
Hide file tree
Showing 7 changed files with 161 additions and 94 deletions.
7 changes: 4 additions & 3 deletions tests/kit/model_zoo/transformers/opt.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,8 @@ def data_gen_for_question_answering():
output_transform_fn = lambda x: x
loss_fn_for_opt_model = lambda x: torch.nn.functional.mse_loss(x.last_hidden_state, torch.ones_like(x.last_hidden_state)
)
loss_fn_for_lm = lambda x: x.loss
loss_fn_for_question_answering = lambda x: x.loss
loss_fn_for_lm = lambda output: output.logits.mean()
config = transformers.OPTConfig(
hidden_size=128,
num_hidden_layers=2,
Expand All @@ -73,11 +74,11 @@ def data_gen_for_question_answering():
model_fn=lambda: transformers.OPTForQuestionAnswering(config),
data_gen_fn=data_gen_for_question_answering,
output_transform_fn=output_transform_fn,
loss_fn=loss_fn_for_lm,
loss_fn=loss_fn_for_question_answering,
model_attribute=ModelAttribute(has_control_flow=True))
model_zoo.register(name='transformers_opt_for_sequence_classification',
model_fn=lambda: transformers.OPTForSequenceClassification(config),
data_gen_fn=data_gen_for_sequence_classification,
data_gen_fn=data_gen,
output_transform_fn=output_transform_fn,
loss_fn=loss_fn_for_lm,
model_attribute=ModelAttribute(has_control_flow=True))
39 changes: 27 additions & 12 deletions tests/test_shardformer/test_model/test_shard_bloom.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,11 +36,14 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,

# check last hidden state & loss
if stage_manager is None or stage_manager.is_last_stage():

if test_config['precision'] == 'fp32':
atol, rtol = 1e-5, 1e-3
else:
atol, rtol = 5e-3, 5e-3
if org_model.__class__.__name__ == 'BloomModel':
check_output_hidden_state(org_output, sharded_output, stage_manager, atol=1e-5, rtol=1e-3)
check_output_hidden_state(org_output, sharded_output, stage_manager, atol=atol, rtol=rtol)

check_loss(org_loss, sharded_loss, atol=1e-6, rtol=1e-3)
check_loss(org_loss, sharded_loss, atol=atol, rtol=rtol)

# unwrap model
if org_model.__class__.__name__ == 'BloomModel':
Expand All @@ -54,14 +57,22 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
row_layer_for_check = ['h[0].self_attention.query_key_value', 'word_embeddings']
col_layer_for_check = ['h[0].self_attention.dense']
if stage_manager is None or stage_manager.is_first_stage():
check_grad(bloom, sharded_bloom, row_layer_for_check, tp_group, atol=1e-6, rtol=1e-5, dim=0, verbose=False)
check_grad(bloom, sharded_bloom, col_layer_for_check, tp_group, atol=1e-6, rtol=1e-5, dim=1, verbose=False)
if test_config['precision'] == 'fp32':
atol, rtol = 1e-6, 1e-5
else:
atol, rtol = 5e-3, 5e-3
check_grad(bloom, sharded_bloom, row_layer_for_check, tp_group, atol=atol, rtol=rtol, dim=0, verbose=False)
check_grad(bloom, sharded_bloom, col_layer_for_check, tp_group, atol=atol, rtol=rtol, dim=1, verbose=False)

# check weights after optimizer.step()
org_optimizer.step()
sharded_optimizer.step()
if stage_manager is None or stage_manager.is_first_stage():
check_weight(bloom, sharded_bloom, col_layer_for_check, tp_group, atol=1e-4, rtol=1e-3, dim=1, verbose=False)
if test_config['precision'] == 'fp32':
atol, rtol = 1e-4, 1e-3
else:
atol, rtol = 5e-3, 5e-3
check_weight(bloom, sharded_bloom, col_layer_for_check, tp_group, atol=atol, rtol=rtol, dim=1, verbose=False)

torch.cuda.empty_cache()

Expand All @@ -70,19 +81,23 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
'tp_size': 2,
'pp_size': 2,
'num_microbatches': 4,
'enable_fused_normalization': True,
'use_lazy_init': True
'enable_all_optimization': True,
'use_lazy_init': True,
'precision': 'fp16',
'initial_scale': 1,
}, {
'tp_size': 1,
'pp_size': 2,
'num_microbatches': 4,
'enable_fused_normalization': False,
'use_lazy_init': False
'enable_all_optimization': False,
'use_lazy_init': False,
'precision': 'fp32',
}, {
'tp_size': 4,
'pp_size': 1,
'enable_fused_normalization': True,
'use_lazy_init': False
'enable_all_optimization': True,
'use_lazy_init': False,
'precision': 'fp32',
}])
def run_bloom_test(test_config):

Expand Down
47 changes: 30 additions & 17 deletions tests/test_shardformer/test_model/test_shard_chatglm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,15 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,

# check last hidden state & loss
if stage_manager is None or stage_manager.is_last_stage():
if test_config['precision'] == 'fp32':
atol, rtol = 1e-5, 1e-3
else:
atol, rtol = 5e-3, 5e-3

if org_model.__class__.__name__ == 'ChatGLMModel':
check_output_hidden_state(org_output, sharded_output, stage_manager, atol=1e-5, rtol=1e-3, dim=1)
check_output_hidden_state(org_output, sharded_output, stage_manager, atol=atol, rtol=rtol, dim=1)

check_loss(org_loss, sharded_loss, atol=1e-6, rtol=1e-3)
check_loss(org_loss, sharded_loss, atol=atol, rtol=rtol)

# unwrap model
if org_model.__class__.__name__ == 'ChatGLMModel':
Expand All @@ -55,34 +59,42 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
row_layer_for_check = ['encoder.layers[0].self_attention.query_key_value', 'embedding.word_embeddings']
col_layer_for_check = ['encoder.layers[0].self_attention.dense']
if stage_manager is None or stage_manager.is_first_stage():
if test_config['precision'] == 'fp32':
atol, rtol = 1e-6, 1e-3
else:
atol, rtol = 5e-3, 5e-3
check_grad(chatglm_model,
shard_chatglm_model,
row_layer_for_check,
tp_group,
atol=1e-6,
rtol=1e-3,
atol=atol,
rtol=rtol,
dim=0,
verbose=False)

check_grad(chatglm_model,
shard_chatglm_model,
col_layer_for_check,
tp_group,
atol=1e-6,
rtol=1e-3,
atol=atol,
rtol=rtol,
dim=1,
verbose=False)

# check weights after optimizer.step()
org_optimizer.step()
sharded_optimizer.step()
if stage_manager is None or stage_manager.is_first_stage():
if test_config['precision'] == 'fp32':
atol, rtol = 1e-4, 1e-3
else:
atol, rtol = 5e-3, 5e-3
check_weight(chatglm_model,
shard_chatglm_model,
col_layer_for_check,
tp_group,
atol=1e-4,
rtol=1e-3,
atol=atol,
rtol=rtol,
dim=1,
verbose=False)

Expand All @@ -93,26 +105,27 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
'tp_size': 2,
'pp_size': 2,
'num_microbatches': 4,
'enable_fused_normalization': True,
'use_lazy_init': True
'enable_all_optimization': True,
'use_lazy_init': True,
'precision': 'fp16',
'initial_scale': 1,
}, {
'tp_size': 1,
'pp_size': 2,
'num_microbatches': 4,
'enable_fused_normalization': False,
'use_lazy_init': False
'enable_all_optimization': False,
'use_lazy_init': False,
'precision': 'fp32',
}, {
'tp_size': 4,
'pp_size': 1,
'enable_fused_normalization': True,
'use_lazy_init': False
'enable_all_optimization': True,
'use_lazy_init': False,
'precision': 'fp32',
}])
def run_chatglm_test(test_config):

# TODO: add test_config for TP+DP after supporting & debugging it
# {'tp_size': 2, 'pp_size': 1, 'enable_fused_normalization': True}

# TODO: add test_config for flash attention & jit operator after supporting

sub_model_zoo = model_zoo.get_sub_registry('transformers_chatglm')
test_config['precision'] = 'float' # Do not use fp16/bf16 in testing
Expand Down
16 changes: 8 additions & 8 deletions tests/test_shardformer/test_model/test_shard_gpt2.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,22 +63,22 @@ def unwrap(module):
row_layer_for_check = ['wte', 'h[0].mlp.c_proj']

# check grad
if test_config['precision'] == 'fp32':
atol, rtol = 1e-4, 1e-3
else:
atol, rtol = 5e-3, 5e-3
if stage_manager is None or stage_manager.is_first_stage():
if test_config['precision'] == 'fp32':
atol, rtol = 1e-4, 1e-3
else:
atol, rtol = 5e-3, 5e-3
check_grad(gpt2, sharded_gpt2, col_layer_for_check, tp_group, atol=atol, rtol=rtol, dim=1, verbose=False)
check_grad(gpt2, sharded_gpt2, row_layer_for_check, tp_group, atol=atol, rtol=rtol, dim=0, verbose=False)

# check weights after optimizer.step()
org_optimizer.step()
sharded_optimizer.step()
if test_config['precision'] == 'fp32':
atol, rtol = 5e-3, 1e-3
else:
atol, rtol = 5e-3, 5e-3
if stage_manager is None or stage_manager.is_first_stage():
if test_config['precision'] == 'fp32':
atol, rtol = 5e-3, 1e-3
else:
atol, rtol = 5e-3, 5e-3
check_weight(gpt2, sharded_gpt2, col_layer_for_check, tp_group, atol=atol, rtol=rtol, dim=1, verbose=False)

torch.cuda.empty_cache()
Expand Down
48 changes: 31 additions & 17 deletions tests/test_shardformer/test_model/test_shard_llama.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,11 +41,15 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,

# check last hidden state & loss
if stage_manager is None or stage_manager.is_last_stage():
if test_config['precision'] == 'fp32':
atol, rtol = 1e-5, 1e-3
else:
atol, rtol = 5e-3, 5e-3

if org_model.__class__.__name__ == 'LlamaModel':
check_output_hidden_state(org_output, sharded_output, stage_manager, atol=1e-5, rtol=1e-3)
check_output_hidden_state(org_output, sharded_output, stage_manager, atol=atol, rtol=rtol)

check_loss(org_loss, sharded_loss, atol=1e-6, rtol=1e-3)
check_loss(org_loss, sharded_loss, atol=atol, rtol=rtol)

# unwrap model
if org_model.__class__.__name__ == 'LlamaModel':
Expand All @@ -59,33 +63,41 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
row_layer_for_check = ['layers[0].self_attn.q_proj', 'embed_tokens']
col_layer_for_check = ['layers[0].self_attn.o_proj']
if stage_manager is None or stage_manager.is_first_stage():
if test_config['precision'] == 'fp32':
atol, rtol = 1e-6, 1e-4
else:
atol, rtol = 5e-3, 5e-3
check_grad(llama_model,
shard_llama_model,
row_layer_for_check,
tp_group,
atol=1e-6,
rtol=1e-4,
atol=atol,
rtol=rtol,
dim=0,
verbose=False)
check_grad(llama_model,
shard_llama_model,
col_layer_for_check,
tp_group,
atol=1e-6,
rtol=1e-4,
atol=atol,
rtol=rtol,
dim=1,
verbose=False)

# check weights after optimizer.step()
org_optimizer.step()
sharded_optimizer.step()
if stage_manager is None or stage_manager.is_first_stage():
if test_config['precision'] == 'fp32':
atol, rtol = 1e-4, 1e-3
else:
atol, rtol = 5e-3, 5e-3
check_weight(llama_model,
shard_llama_model,
col_layer_for_check,
tp_group,
atol=1e-4,
rtol=1e-3,
atol=atol,
rtol=rtol,
dim=1,
verbose=False)

Expand All @@ -96,30 +108,32 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
'tp_size': 2,
'pp_size': 2,
'num_microbatches': 2,
'enable_fused_normalization': True,
'use_lazy_init': True
'enable_all_optimization': True,
'use_lazy_init': True,
'precision': 'fp16',
'initial_scale': 1,
}, {
'tp_size': 1,
'pp_size': 2,
'num_microbatches': 4,
'use_lazy_init': False
'use_lazy_init': False,
'precision': 'fp32',
}, {
'tp_size': 4,
'pp_size': 1,
'enable_fused_normalization': True,
'use_lazy_init': False
'enable_all_optimization': True,
'use_lazy_init': False,
'precision': 'fp32',
}, {
'tp_size': 1,
'pp_size': 4,
'num_microbatches': 4,
'use_lazy_init': False
'use_lazy_init': False,
'precision': 'fp32',
}])
def run_llama_test(test_config):

# TODO: add test_config for TP+DP after supporting & debugging it
# {'tp_size': 2, 'pp_size': 1, 'enable_fused_normalization': True}

# TODO: add test_config for flash attention & jit operator after supporting

sub_model_zoo = model_zoo.get_sub_registry('transformers_llama')
test_config['precision'] = 'float' # Do not use fp16/bf16 in testing
Expand Down
Loading

0 comments on commit 759646c

Please sign in to comment.