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) reinitialize_optimizer(optimizer, model)
if self.zero_stage == 0: if self.zero_stage == 0:
assert self.ep_size > 1 # assert self.ep_size > 1
if self.precision in ["fp16", "bf16"]: if self.precision in ["fp16", "bf16"]:
optimizer = HybridParallelAMPOptimizer( 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 output_transform_fn = lambda x: x
# define loss function # define loss function
loss_fn_for_mixtral_model = lambda x: torch.nn.functional.mse_loss( loss_fn_for_mixtral_model = lambda x: torch.nn.functional.mse_loss(x[0], torch.ones_like(x[0]))
x.last_hidden_state, torch.ones_like(x.last_hidden_state)
)
loss_fn = lambda x: x.loss loss_fn = lambda x: x.loss
loss_fn_for_seq_classification = lambda output: output.logits.mean() loss_fn_for_seq_classification = lambda output: output.logits.mean()
config = MixtralConfig( 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"): if hasattr(config, "pad_token_id"):
@ -64,19 +67,19 @@ model_zoo.register(
loss_fn=loss_fn_for_mixtral_model, loss_fn=loss_fn_for_mixtral_model,
model_attribute=ModelAttribute(has_control_flow=True), model_attribute=ModelAttribute(has_control_flow=True),
) )
model_zoo.register( # model_zoo.register(
name="transformers_mixtral_for_casual_lm", # name="transformers_mixtral_for_casual_lm",
model_fn=lambda: transformers.MixtralForCausalLM(config), # model_fn=lambda: transformers.MixtralForCausalLM(config),
data_gen_fn=data_gen_for_lm, # data_gen_fn=data_gen_for_lm,
output_transform_fn=output_transform_fn, # output_transform_fn=output_transform_fn,
loss_fn=loss_fn, # loss_fn=loss_fn,
model_attribute=ModelAttribute(has_control_flow=True), # model_attribute=ModelAttribute(has_control_flow=True),
) # )
model_zoo.register( # model_zoo.register(
name="transformers_mixtral_for_sequence_classification", # name="transformers_mixtral_for_sequence_classification",
model_fn=lambda: transformers.MixtralForSequenceClassification(config), # model_fn=lambda: transformers.MixtralForSequenceClassification(config),
data_gen_fn=data_gen_for_sequence_classification, # data_gen_fn=data_gen_for_sequence_classification,
output_transform_fn=output_transform_fn, # output_transform_fn=output_transform_fn,
loss_fn=loss_fn_for_seq_classification, # loss_fn=loss_fn_for_seq_classification,
model_attribute=ModelAttribute(has_control_flow=True), # 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, "tp_size": 1,
"pp_size": 1, "pp_size": 4,
"ep_size": 4, "ep_size": 1,
"num_microbatches": 2, "num_microbatches": 4,
"zero_stage": 0, "zero_stage": 0,
"enable_all_optimization": True, "enable_all_optimization": True,
"use_lazy_init": False, "use_lazy_init": False,
"precision": "fp16", "precision": "fp16",
"initial_scale": 1, "initial_scale": 1,
}, },
{ # {
"tp_size": 1, # "tp_size": 1,
"pp_size": 1, # "pp_size": 1,
"ep_size": 4, # "ep_size": 4,
"num_microbatches": 2, # "num_microbatches": 2,
"zero_stage": 1, # "zero_stage": 1,
"enable_all_optimization": True, # "enable_all_optimization": True,
"use_lazy_init": False, # "use_lazy_init": False,
"precision": "fp16", # "precision": "fp16",
"initial_scale": 1, # "initial_scale": 1,
}, # },
{ # {
"tp_size": 1, # "tp_size": 1,
"pp_size": 1, # "pp_size": 1,
"ep_size": 4, # "ep_size": 4,
"num_microbatches": 2, # "num_microbatches": 2,
"zero_stage": 2, # "zero_stage": 2,
"enable_all_optimization": True, # "enable_all_optimization": True,
"use_lazy_init": False, # "use_lazy_init": False,
"precision": "fp16", # "precision": "fp16",
"initial_scale": 1, # "initial_scale": 1,
}, # },
], ],
) )
def run_mixtral_test(test_config): def run_mixtral_test(test_config):