github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[test] mixtra pp shard test

moe_sp
hxwang 2024-07-04 06:39:01 +00:00
parent 5a9490a46b
commit 4fc6f9aa98
No known key found for this signature in database
GPG Key ID: 0EC383D418F0B9F8
3 changed files with 49 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/booster/plugin/moe_hybrid_parallel_plugin.py
View File

@ -411,7 +411,7 @@ class MoeHybridParallelPlugin(HybridParallelPlugin):
reinitialize_optimizer(optimizer, model)
if self.zero_stage == 0:
assert self.ep_size > 1
# assert self.ep_size > 1
if self.precision in ["fp16", "bf16"]:
optimizer = HybridParallelAMPOptimizer(

43
tests/kit/model_zoo/transformers/mixtral.py
View File

@ -43,14 +43,17 @@ def data_gen_for_sequence_classification():
output_transform_fn = lambda x: x
# define loss function
loss_fn_for_mixtral_model = lambda x: torch.nn.functional.mse_loss(
x.last_hidden_state, torch.ones_like(x.last_hidden_state)
)
loss_fn_for_mixtral_model = lambda x: torch.nn.functional.mse_loss(x[0], torch.ones_like(x[0]))
loss_fn = lambda x: x.loss
loss_fn_for_seq_classification = lambda output: output.logits.mean()
config = MixtralConfig(
hidden_size=256, intermediate_size=256, num_attention_heads=64, num_hidden_layers=2, vocab_size=50258
hidden_size=256,
intermediate_size=256,
num_attention_heads=64,
num_hidden_layers=2,
vocab_size=50258,
output_router_logits=True,
)
if hasattr(config, "pad_token_id"):
@ -64,19 +67,19 @@ model_zoo.register(
loss_fn=loss_fn_for_mixtral_model,
model_attribute=ModelAttribute(has_control_flow=True),
)
model_zoo.register(
name="transformers_mixtral_for_casual_lm",
model_fn=lambda: transformers.MixtralForCausalLM(config),
data_gen_fn=data_gen_for_lm,
output_transform_fn=output_transform_fn,
loss_fn=loss_fn,
model_attribute=ModelAttribute(has_control_flow=True),
)
model_zoo.register(
name="transformers_mixtral_for_sequence_classification",
model_fn=lambda: transformers.MixtralForSequenceClassification(config),
data_gen_fn=data_gen_for_sequence_classification,
output_transform_fn=output_transform_fn,
loss_fn=loss_fn_for_seq_classification,
model_attribute=ModelAttribute(has_control_flow=True),
)
# model_zoo.register(
# name="transformers_mixtral_for_casual_lm",
# model_fn=lambda: transformers.MixtralForCausalLM(config),
# data_gen_fn=data_gen_for_lm,
# output_transform_fn=output_transform_fn,
# loss_fn=loss_fn,
# model_attribute=ModelAttribute(has_control_flow=True),
# )
# model_zoo.register(
# name="transformers_mixtral_for_sequence_classification",
# model_fn=lambda: transformers.MixtralForSequenceClassification(config),
# data_gen_fn=data_gen_for_sequence_classification,
# output_transform_fn=output_transform_fn,
# loss_fn=loss_fn_for_seq_classification,
# model_attribute=ModelAttribute(has_control_flow=True),
# )

50
tests/test_shardformer/test_model/test_shard_mixtral.py
View File

@ -114,37 +114,37 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
[
{
"tp_size": 1,
"pp_size": 1,
"ep_size": 4,
"num_microbatches": 2,
"pp_size": 4,
"ep_size": 1,
"num_microbatches": 4,
"zero_stage": 0,
"enable_all_optimization": True,
"use_lazy_init": False,
"precision": "fp16",
"initial_scale": 1,
},
{
"tp_size": 1,
"pp_size": 1,
"ep_size": 4,
"num_microbatches": 2,
"zero_stage": 1,
"enable_all_optimization": True,
"use_lazy_init": False,
"precision": "fp16",
"initial_scale": 1,
},
{
"tp_size": 1,
"pp_size": 1,
"ep_size": 4,
"num_microbatches": 2,
"zero_stage": 2,
"enable_all_optimization": True,
"use_lazy_init": False,
"precision": "fp16",
"initial_scale": 1,
},
# {
# "tp_size": 1,
# "pp_size": 1,
# "ep_size": 4,
# "num_microbatches": 2,
# "zero_stage": 1,
# "enable_all_optimization": True,
# "use_lazy_init": False,
# "precision": "fp16",
# "initial_scale": 1,
# },
# {
# "tp_size": 1,
# "pp_size": 1,
# "ep_size": 4,
# "num_microbatches": 2,
# "zero_stage": 2,
# "enable_all_optimization": True,
# "use_lazy_init": False,
# "precision": "fp16",
# "initial_scale": 1,
# },
],
)
def run_mixtral_test(test_config):