allenai / allenai/open-instruct

olmo-core runs don't record their Beaker experiment in wandb

未关闭
#1,812 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看
主要语言
Python
星标
3.9k
派生
585
平均合并
5 天 17 小时
30 天内合并 PR
16

描述

`olmo_core_finetune.py` / `dpo.py` runs log to wandb with no reference to the Beaker job that produced them, so there is no way to get from a wandb run to its experiment (or back) without grepping job logs and matching timestamps.

## Already solved elsewhere

`utils.maybe_get_beaker_config()` returns a `BeakerRuntimeConfig` with `beaker_workload_id`, `beaker_experiment_url`, `beaker_node_hostname` and dataset ids. Both non-olmo-core paths merge it into the wandb config:

- `grpo_fast.py:1062` — `all_configs.update(vars(beaker_config))`
- `dpo_tune_cache.py:198` — `experiment_config.update(vars(beaker_config))`

The olmo-core path never calls it. `olmo_core_utils.build_common_callbacks` passes `config=config_dict` straight to `WandBCallback`, and that dict has no Beaker fields.

## Suggested fix

Merge `maybe_get_beaker_config()` into the config passed to `WandBCallback` in `build_common_callbacks`, matching what the other two paths do. Guarded by `is_beaker_job()`, so local runs are unaffected. This also covers the olmo-core DPO path, which shares those callbacks.

## Why it matters

Tracing a finished 7B SFT run back to its Beaker job currently requires `beaker experiment logs | grep 'run--'`, or matching the wandb `host` and `startedAt` against job metadata. The wandb run's `root` path does contain the Beaker-minted `$CHECKPOINT_OUTPUT_DIR` suffix, but that lives under `deletable_checkpoint_states` and is not durable — and it is ambiguous when a run resumes into a previous run's folder, where one directory maps to two experiments.

Concrete example: wandb run `ed453531` in `ai2-llm/open_instruct_internal` is Beaker experiment `01KZHT32T30M2VHCWKRJS1G9P7`, and confirming that took log grepping plus a timestamp match.

贡献指南

打开贡献指南

评估

这个 Issue 还没有评估数据。

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。