Initialize wandb experiment.
()
| 312 | |
| 313 | |
| 314 | def initialize_wandb_experiment(): |
| 315 | """Initialize wandb experiment.""" |
| 316 | assert wandb is not None, "Fail to import wandb" |
| 317 | |
| 318 | args = get_args() |
| 319 | config = args.__dict__ |
| 320 | |
| 321 | wandb_id_path = os.path.join(args.save, "wandb_id.txt") |
| 322 | if os.path.exists(wandb_id_path): |
| 323 | wandb_id = open(wandb_id_path, "r").read().strip() |
| 324 | else: |
| 325 | wandb_id = wandb.util.generate_id() |
| 326 | open(wandb_id_path, "w").write(wandb_id) |
| 327 | |
| 328 | wandb.init(id=wandb_id, project="megatron", config=config, resume="allow") |
| 329 | |
| 330 | |
| 331 | def _initialize_mem_buffs(): |