tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2025-12-06 04:12:13 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Replace wandb initialization with get_wandb function

Browse Source
This commit is contained in:
Sermet Pekin 2025-11-05 15:56:49 +03:00 committed by GitHub
parent 679ac96efe
commit 523714b5c8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
scripts/chat_rl.py
View File

@ -19,11 +19,12 @@ torchrun --standalone --nproc_per_node=8 -m scripts.chat_rl -- --run=default
import os
import itertools
import re
import wandb
import torch
import torch.distributed as dist
from nanochat.common import compute_init, compute_cleanup, print0, get_base_dir, DummyWandb
from nanochat.common import compute_init, compute_cleanup, print0, get_base_dir, get_wandb
from nanochat.checkpoint_manager import save_checkpoint, load_model
from nanochat.engine import Engine
from tasks.gsm8k import GSM8K
@ -60,8 +61,7 @@ dtype = torch.float32 if dtype == 'float32' else torch.bfloat16
autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=dtype)
# wandb logging init
use_dummy_wandb = run == "dummy" or not master_process
wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-rl", name=run, config=user_config)
wandb_run = get_wandb("nanochat-rl" , run=run, master_process=master_process, user_config=user_config)
# Init model and tokenizer
model, tokenizer, meta = load_model(source, device, phase="eval")