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

[hotfix] fix typo of openmoe model source (#5403)

pull/5335/head^2
Luo Yihang 2024-03-05 21:44:38 +08:00 committed by GitHub
parent e304e4db35
commit e239cf9060
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 7 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
examples/language/openmoe/benchmark/benchmark_cai.py
View File

@ -207,7 +207,7 @@ def main():
coordinator.print_on_master(f"Set plugin as {plugin}") coordinator.print_on_master(f"Set plugin as {plugin}")
# Build OpenMoe model # Build OpenMoe model
repo_name = "hpcaitech/openmoe-" + args.model_name repo_name = "hpcai-tech/openmoe-" + args.model_name
config = LlamaConfig.from_pretrained(repo_name) config = LlamaConfig.from_pretrained(repo_name)
set_openmoe_args( set_openmoe_args(
config, config,

2
examples/language/openmoe/benchmark/benchmark_fsdp.py
View File

@ -53,7 +53,7 @@ def fsdp_main(rank, world_size, args):
train_loader = torch.utils.data.DataLoader(dataset, **train_kwargs) train_loader = torch.utils.data.DataLoader(dataset, **train_kwargs)
torch.cuda.set_device(rank) torch.cuda.set_device(rank)
config = LlamaConfig.from_pretrained("hpcaitech/openmoe-%s" % args.model_name) config = LlamaConfig.from_pretrained("hpcai-tech/openmoe-%s" % args.model_name)
set_openmoe_args( set_openmoe_args(
config, config,
num_experts=config.num_experts, num_experts=config.num_experts,

6
examples/language/openmoe/infer.py
View File

@ -15,19 +15,19 @@ def parse_args():
def inference(args): def inference(args):
tokenizer = T5Tokenizer.from_pretrained("google/umt5-small") tokenizer = T5Tokenizer.from_pretrained("google/umt5-small")
if args.model == "test": if args.model == "test":
config = LlamaConfig.from_pretrained("hpcaitech/openmoe-base") config = LlamaConfig.from_pretrained("hpcai-tech/openmoe-base")
set_openmoe_args(config, set_openmoe_args(config,
num_experts=config.num_experts, num_experts=config.num_experts,
moe_layer_interval=config.moe_layer_interval, moe_layer_interval=config.moe_layer_interval,
enable_kernel=True) enable_kernel=True)
model = OpenMoeForCausalLM(config) model = OpenMoeForCausalLM(config)
else: else:
config = LlamaConfig.from_pretrained(f"hpcaitech/openmoe-{args.model}") config = LlamaConfig.from_pretrained(f"hpcai-tech/openmoe-{args.model}")
set_openmoe_args(config, set_openmoe_args(config,
num_experts=config.num_experts, num_experts=config.num_experts,
moe_layer_interval=config.moe_layer_interval, moe_layer_interval=config.moe_layer_interval,
enable_kernel=False) enable_kernel=False)
model = OpenMoeForCausalLM.from_pretrained(f"hpcaitech/openmoe-{args.model}", config=config) model = OpenMoeForCausalLM.from_pretrained(f"hpcai-tech/openmoe-{args.model}", config=config)
model = model.eval().bfloat16() model = model.eval().bfloat16()
model = model.to(torch.cuda.current_device()) model = model.to(torch.cuda.current_device())

4
examples/language/openmoe/train.py
View File

@ -269,12 +269,12 @@ def main():
# Build OpenMoe model # Build OpenMoe model
if test_mode: if test_mode:
config = LlamaConfig.from_pretrained("hpcaitech/openmoe-base") config = LlamaConfig.from_pretrained("hpcai-tech/openmoe-base")
config.hidden_size = 128 config.hidden_size = 128
config.intermediate_size = 256 config.intermediate_size = 256
config.vocab_size = 32000 config.vocab_size = 32000
else: else:
repo_name = "hpcaitech/openmoe-" + args.model_name repo_name = "hpcai-tech/openmoe-" + args.model_name
config = LlamaConfig.from_pretrained(repo_name) config = LlamaConfig.from_pretrained(repo_name)
set_openmoe_args( set_openmoe_args(
config, config,