squall/picotron
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

small fix on world_size with pgm

Browse Source
This commit is contained in:
ferdinand.mom 2024-12-02 18:12:02 +00:00
parent b6267c768e
commit 75939867d9
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
train.py
View File

@ -126,10 +126,10 @@ if __name__ == "__main__":
project="picotron",
name=f"{config['logging']['run_name']}_{tokens_per_step}_{pgm.process_group_manager}",
config={
"tensor_parallel_size": pgm.process_group_manager.tp_size,
"context_parallel_size": pgm.process_group_manager.cp_size,
"pipeline_parallel_size": pgm.process_group_manager.pp_size,
"data_parallel_size": pgm.process_group_manager.dp_size,
"tensor_parallel_size": pgm.process_group_manager.tp_world_size,
"context_parallel_size": pgm.process_group_manager.cp_world_size,
"pipeline_parallel_size": pgm.process_group_manager.pp_world_size,
"data_parallel_size": pgm.process_group_manager.dp_world_size,
"model": config["model"]["name"],
"dataset": config["dataset"]["name"],
"max_tokens": config["training"]["max_tokens"],