hailin
/
jd_train
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

.

This commit is contained in:
hailin 2025-09-04 18:52:07 +08:00
parent b27faa397c
commit 905d5c4542
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
train_sft_ds.py
View File

@ -4,10 +4,7 @@ import os
os.environ.pop("PYTHONNOUSERSITE", None) os.environ.pop("PYTHONNOUSERSITE", None)
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
# ✅ 只有 rank0 用 wandb，其它 rank 不上报
if os.environ.get("RANK", "0") != "0" and args.report_to == "wandb":
print(f"[rank {os.environ.get('RANK')}] force report_to=none", flush=True)
args.report_to = "none"
os.environ.setdefault("WANDB_START_METHOD", "thread") os.environ.setdefault("WANDB_START_METHOD", "thread")
os.environ.setdefault("WANDB_DIR", f"/tmp/{os.environ.get('USER','user')}/wandb") os.environ.setdefault("WANDB_DIR", f"/tmp/{os.environ.get('USER','user')}/wandb")
@ -451,6 +448,12 @@ def parse_args():
def main(): def main():
args = parse_args() args = parse_args()
# ✅ 只有 rank0 用 wandb，其它 rank 不上报
if os.environ.get("RANK", "0") != "0" and args.report_to == "wandb":
print(f"[rank {os.environ.get('RANK')}] force report_to=none", flush=True)
args.report_to = "none"
set_seed(args.seed) set_seed(args.seed)
host = socket.gethostname() host = socket.gethostname()