Skip to content

Fix get_wandb_tags tag-limit handling and log all GRPO configs #1727

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mnoukhov wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Fix get_wandb_tags tag-limit handling and log all GRPO configs #1727

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c4de080
correct wandb tag limit
mnoukhov Jun 17, 2026
dd15db7
log all args
mnoukhov Jun 15, 2026
bf5e0f9
Update CHANGELOG
mnoukhov Jun 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,3 +15,5 @@ All notable changes to this project will be documented in this file.
### Removed

### Fixed
- Make `get_wandb_tags` accept `extra_tags` (e.g. the experiment name) and truncate them to W&B's 64-char tag limit, and pass tags to the OLMo-core GRPO `WandBCallback` (https://github.com/allenai/open-instruct/pull/1727).
- Log all config dataclasses (`tc`, `model_config`, `streaming_config`, `vllm_config`) in the OLMo-core GRPO wandb config, not just `args` (https://github.com/allenai/open-instruct/pull/1727).
2 changes: 1 addition & 1 deletion open_instruct/dpo_tune_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ def main(args: dpo_utils.DPOExperimentConfig, tc: TokenizerConfig):
"wandb": {
"name": args.exp_name,
"entity": args.wandb_entity,
"tags": [args.exp_name] + get_wandb_tags(),
"tags": get_wandb_tags(extra_tags=[args.exp_name]),
}
},
)
Expand Down
2 changes: 1 addition & 1 deletion open_instruct/finetune.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -462,7 +462,7 @@ def main(args: FlatArguments, tc: TokenizerConfig):
"wandb": {
"name": args.exp_name,
"entity": args.wandb_entity,
"tags": [args.exp_name] + get_wandb_tags(),
"tags": get_wandb_tags(extra_tags=[args.exp_name]),
}
},
)
Expand Down
9 changes: 8 additions & 1 deletion open_instruct/grpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -296,9 +296,16 @@ def main(
)
logger.info("======== Model update group setup successfully =========")

json_config = dataclasses.asdict(args)
json_config = {}
if beaker_config is not None:
json_config.update(dataclasses.asdict(beaker_config))
json_config.update(
**dataclasses.asdict(args),
**dataclasses.asdict(tc),
**dataclasses.asdict(model_config),
**dataclasses.asdict(streaming_config),
**dataclasses.asdict(vllm_config),
)
Comment on lines +302 to +308

@gemini-code-assist gemini-code-assist Bot Jun 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Unpacking multiple dataclasses as keyword arguments using ** in a single update() call is risky. If any of these dataclasses (args, tc, model_config, streaming_config, vllm_config) share any field names (either now or in the future), Python will raise a TypeError: update() got multiple values for keyword argument ... at runtime.

To prevent potential runtime crashes and ensure robustness, update the dictionary sequentially by passing the dictionaries directly to .update().

    for config in (args, tc, model_config, streaming_config, vllm_config):
        json_config.update(dataclasses.asdict(config))

@mnoukhov mnoukhov Jun 17, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would silently overwrite some configs with others. I prefer to have it fail if two dataclasses share configs

utils.ray_get_with_progress(
[
m.setup_callbacks.remote(
Expand Down
2 changes: 1 addition & 1 deletion open_instruct/grpo_fast.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1072,7 +1072,7 @@ def setup_experiment_tracking(
group=args.wandb_group_name,
name=args.run_name,
save_code=True,
tags=[args.exp_name] + get_wandb_tags(),
tags=get_wandb_tags(extra_tags=[args.exp_name]),
)
# Set training_step as the default x-axis metric
wandb.define_metric("training_step")
Expand Down
6 changes: 5 additions & 1 deletion open_instruct/grpo_olmo_core_actor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -357,7 +357,11 @@ def fit(self) -> dict:

if self.with_tracking:
trainer_callbacks["wandb"] = callbacks.WandBCallback(
name=self.run_name, project=self.wandb_project, entity=self.wandb_entity, config=self.json_config
name=self.run_name,
project=self.wandb_project,
entity=self.wandb_entity,
config=self.json_config,
tags=utils.get_wandb_tags(extra_tags=[self.grpo_config.exp_name]),
)

if self.rank == 0 and self.eval_dataset is not None and self.grpo_config.local_eval_every > 0:
Expand Down
2 changes: 1 addition & 1 deletion open_instruct/reward_modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,7 @@ def main(args: Args, tc: TokenizerConfig, model_config: ModelConfig):
config=all_configs,
name=args.run_name,
save_code=True,
tags=[args.exp_name] + get_wandb_tags(),
tags=get_wandb_tags(extra_tags=[args.exp_name]),
)
writer = SummaryWriter(f"runs/{args.run_name}")
hyperparams_table = "\n".join([f"|{key}|{value}|" for key, value in vars(args).items()])
Expand Down
10 changes: 7 additions & 3 deletions open_instruct/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -764,9 +764,13 @@ def get_git_commit() -> str:
return os.environ.get("GIT_COMMIT", "unknown")


def get_wandb_tags() -> list[str]:
"""Get tags for Weights & Biases (e.g., `no-tag-404-g98dc659,pr-123,branch-main`)"""
tags = [t for t in os.environ.get("WANDB_TAGS", "").split(",") if t != ""]
def get_wandb_tags(extra_tags: list[str] | None = None) -> list[str]:
"""Get tags for Weights & Biases (e.g., `no-tag-404-g98dc659,pr-123,branch-main`).

`extra_tags` (e.g. the experiment name) are prepended and, like all tags, truncated to W&B's 64-char limit.
"""
tags = list(extra_tags) if extra_tags else []
tags += [t for t in os.environ.get("WANDB_TAGS", "").split(",") if t != ""]
if (git_commit := get_git_commit()) != "unknown":
tags.append(f"commit: {git_commit}")
try:
Expand Down
Loading