THChou1220 commited on
Commit
ecd8c8c
·
verified ·
1 Parent(s): b4abbf6

Delete upload_trainer_state_to_wandb.py

Browse files
Files changed (1) hide show
  1. upload_trainer_state_to_wandb.py +0 -31
upload_trainer_state_to_wandb.py DELETED
@@ -1,31 +0,0 @@
1
- import json
2
- import wandb
3
-
4
- state_path = "trainer_state.json"
5
-
6
- with open(state_path, "r") as f:
7
- state = json.load(f)
8
-
9
- run = wandb.init(
10
- project="gemma4-finetune",
11
- name="gemma4_e4b_action_stage1_from_trainer_state",
12
- )
13
-
14
- for log in state["log_history"]:
15
- if "loss" in log:
16
- wandb.log({
17
- "train/loss": log["loss"],
18
- "train/learning_rate": log.get("learning_rate"),
19
- "train/grad_norm": log.get("grad_norm"),
20
- "epoch": log.get("epoch"),
21
- }, step=log["step"])
22
-
23
- if "train_loss" in log:
24
- wandb.log({
25
- "train/final_train_loss": log["train_loss"],
26
- "train/runtime": log.get("train_runtime"),
27
- "train/samples_per_second": log.get("train_samples_per_second"),
28
- "train/steps_per_second": log.get("train_steps_per_second"),
29
- }, step=log["step"])
30
-
31
- run.finish()