feat: Refactor training framework with modular architecture #3
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
zzhfz
assigned
baominghelly and
Chamberlain0w0 and unassigned
baominghelly and
Chamberlain0w0
zzhfz
requested review from
Chamberlain0w0 and
baominghelly
and removed request for
Chamberlain0w0 and
baominghelly
baominghelly
requested review from
Chamberlain0w0 and
baominghelly
and removed request for
Chamberlain0w0 and
baominghelly
zzhfz
requested review from
baominghelly
and removed request for
Chamberlain0w0
baominghelly
requested review from
Chamberlain0w0 and
baominghelly
and removed request for
Chamberlain0w0 and
baominghelly
|
|
||
| dp = parallel.get("dp", 1) | ||
| tp = parallel.get("tp", 1) | ||
| pp = parallel.get("pp", {}).get("value", 1) |
Collaborator
There was a problem hiding this comment.
如果pp的类型不是 G-pipe，应该怎么办呢？
| ] | ||
|
|
||
| megatron_args = [ | ||
| "pretrain_gpt.py", |
| torchrun_cmd = [ | ||
| "torchrun", | ||
| f"--nproc_per_node={nproc_per_node}", | ||
| "--master_port=29501" |
| f"--num-attention-heads={num_attention_heads}", | ||
| f"--max-position-embeddings={max_position_embeddings}", | ||
| f"--vocab-size={vocab_size}", | ||
| "--mock-data", |
Collaborator
There was a problem hiding this comment.
如果我们的config里面有数据集，你这里传mock data是不是不行啊？
|
|
||
| # regex patterns | ||
| loss_pattern = re.compile(r"lm loss:\s*([+\-]?\d+(?:\.\d+)?(?:[Ee][+\-]?\d+)?)", re.IGNORECASE) | ||
| #ppl_pattern_alt = re.compile(r"lm loss PPL:\s*([+\-]?\d+(?:\.\d+)?(?:[Ee][+\-]?\d+)?)", re.IGNORECASE) |
| if me: | ||
| try: | ||
| elapsed_ms = float(me.group(1)) | ||
| tokens_per_iter = mbs * seq |
| flog.write(line) | ||
|
|
||
| # try match loss | ||
| m = loss_pattern.search(line) |
Collaborator
There was a problem hiding this comment.
这个地方有没有跳过那些warmup的训练环节呢？
zzhfz
requested review from
Chamberlain0w0 and
baominghelly
and removed request for
baominghelly
zzhfz
changed the title
baominghelly
reviewed
Dec 2, 2025
| } | ||
|
|
||
| # Start training process | ||
| print("Launching:", " ".join(cmd)) |
Collaborator
There was a problem hiding this comment.
We should not use print for logging, we should use logger module here.
Please apply to all code in the pr.
baominghelly
reviewed
Dec 2, 2025
|
|
||
| # Add data configuration | ||
| if self.config.train_dataset is None or (isinstance(self.config.train_dataset, str) and self.config.train_dataset.lower() == "mock"): | ||
| megatron_args += ["--mock-data", "--tokenizer-type", "NullTokenizer", "--vocab-size", str(self.config.vocab_size)] |
Collaborator
There was a problem hiding this comment.
Do we need to use different tokenizer for different model here?
baominghelly
reviewed
Dec 2, 2025
| f"--micro-batch-size={self.config.mbs}", | ||
| f"--global-batch-size={self.config.gbs}", | ||
| f"--seq-length={self.config.seq_len}", | ||
| f"--lr={self.config.lr}", |
Collaborator
There was a problem hiding this comment.
Please also consider decay lr here.
We could add these configs in our json file:
f"--lr_scheduler_type=cosine", # 推荐：余弦退火 (cosine) 或 线性 (linear)
f"--warmup_ratio=0.03", # 推荐：前 3% 的步数用于热身
# 或者使用步数
# f"--warmup_steps=100",
baominghelly
reviewed
Dec 2, 2025
| # Add common parameters | ||
| megatron_args += [ | ||
| "--transformer-impl", "local", | ||
| "--bf16", |
Collaborator
There was a problem hiding this comment.
we should get precision from our json file
baominghelly
reviewed
Dec 2, 2025
| def main(): | ||
| parser = argparse.ArgumentParser() | ||
| parser.add_argument("--config", required=True, help="path to config.json") | ||
| parser.add_argument("--framework", default="megatron", choices=["megatron", "infinitrain"], |
Collaborator
There was a problem hiding this comment.
Could we get this config from json file?
baominghelly
reviewed
Dec 2, 2025
| parser.add_argument("--config", required=True, help="path to config.json") | ||
| parser.add_argument("--framework", default="megatron", choices=["megatron", "infinitrain"], | ||
| help="training framework to use") | ||
| parser.add_argument("--gpu-platform", default="nvidia", choices=["nvidia", "other"], |
Collaborator
There was a problem hiding this comment.
Please add a device arg in config file
baominghelly
requested changes
Dec 2, 2025
Collaborator
baominghelly
left a comment
baominghelly
left a comment
There was a problem hiding this comment.
comment added in pr
zzhfz
requested review from
baominghelly
and removed request for
Chamberlain0w0
zzhfz
commented
Dec 3, 2025
Contributor
Author
zzhfz
left a comment
zzhfz
left a comment
There was a problem hiding this comment.
@baominghelly 训练脚本修改已完成，请审核最新版本。
主要更新
- 标准化输出格式：run_id/testcase符合规范
- 完整配置支持
- 已解决6条评审意见
- 成功通过Megatron-LM测试
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
description
Add Modular Architecture Refactoring script
evidence
./train/train.gpt.946e7e31-adaa-4421-bef1-63eb7726402f_result.json
./train/train.gpt.946e7e31-adaa-4421-bef1-63eb7726402f_train_loss.csv
./train/train.gpt.946e7e31-adaa-4421-bef1-63eb7726402f_train_ppl.csv
./train/train.gpt.946e7e31-adaa-4421-bef1-63eb7726402f_train_throughput.csv