Update dependencies for OLMo-core trainer

[finbarrtimbers]

Summary

• Update ai2-olmo-core to use finbarr/skip-init-if-initialized fork (needed for OLMo-core trainer)
• Add torchao>=0.11.0 for OLMo-core compatibility
• Update vllm from ==0.12.0 to >=0.12.0
• Update wandb to 0.18.1
• Switch PyTorch index from cu129 to cu128

Test plan

• Run make style && make quality
• Run uv run pytest open_instruct/test_utils.py
• Run full test suite on CI

🤖 Generated with Claude Code

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.

[gemini-code-assist]

Summary of Changes

Hello @finbarrtimbers, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request primarily focuses on updating and streamlining the project's Python dependencies and build configurations. The changes are crucial for ensuring compatibility with the OLMo-core trainer, introducing new libraries like torchao, and standardizing the PyTorch CUDA version to cu128. These adjustments aim to resolve potential dependency conflicts and improve the overall stability and maintainability of the development environment.

Highlights

• Dependency Update: ai2-olmo-core: The ai2-olmo-core dependency has been updated from a fixed version (==2.3.0) to a specific Git fork (finbarr/skip-init-if-initialized) to ensure compatibility with the OLMo-core trainer.
• New Dependency: torchao: The torchao library has been added as a new dependency, with a version constraint of >=0.11.0 for non-Darwin systems.
• vllm Version and Platform Simplification: The vllm dependency's version constraint has been relaxed to >=0.12.0, and its platform-specific configurations, including a dedicated aarch64 wheel source, have been removed and simplified.
• wandb Version Pin: The wandb dependency has been explicitly pinned to version 0.18.1.
• PyTorch CUDA Index Switch: The PyTorch index used for installations has been switched from cu129 (and cu130 for aarch64) to cu128 for non-Darwin systems, standardizing the CUDA version.
• Environment Configuration Streamlining: The project's environment configurations have been streamlined by removing specific aarch64 environment definitions.
• Type Checker Source List Adjustment: Several files, including actor_manager.py, benchmark_generators.py, code_utils/api.py, data_loader.py, and dpo_config.py, have been removed from the tool.ty.src include list in pyproject.toml.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request updates several dependencies to support the OLMo-core trainer, including switching to a fork of ai2-olmo-core, adding torchao, and updating vllm and wandb. It also changes the PyTorch CUDA index from cu129 to cu128. My review focuses on potential inconsistencies introduced by these changes, such as a wandb downgrade, conflicting information about flash-attn support on aarch64, outdated documentation, and the removal of several files from static type checking. These changes could impact stability, maintainability, and correctness.

I am having trouble creating individual review comments. Click here to see my feedback.

pyproject.toml (35)

This change removes the platform_machine != 'aarch64' condition, which enables flash-attn for aarch64 architectures. However, the docs/DGX_SPARK.md file states that "Flash Attention: Not available - uses PyTorch SDPA instead" for DGX Spark (aarch64). This seems contradictory. If flash-attn is not supported on aarch64, this change could cause installation or runtime errors on that platform. Please clarify if flash-attn is now supported on aarch64 or if this condition should be kept.

pyproject.toml (32)

The wandb dependency is being downgraded from 0.23.1 to 0.18.1. This is a significant version change. While the PR description mentions updating to 0.18.1, a downgrade of this magnitude could introduce breaking changes or remove important features/bugfixes. Could you please provide some context for this downgrade? If it's intentional, it would be helpful to add a comment in pyproject.toml explaining why this specific older version is required.

pyproject.toml (60-65)

This pull request removes the pinned vLLM wheel for aarch64 that was previously defined here. However, the documentation in docs/DGX_SPARK.md still refers to it:

1. vLLM: Uses cu130 wheel (v0.13.0) pinned to specific commit

The documentation should be updated to reflect that vLLM is no longer pinned to a specific wheel for aarch64 and what the new installation process or behavior is. Keeping documentation in sync with code changes is crucial for maintainability.

pyproject.toml (169-176)

Several files (actor_manager.py, benchmark_generators.py, code_utils/api.py, data_loader.py, dpo_config.py) have been removed from the [tool.ty.src].include list for type checking. Excluding files from static analysis can hide potential type-related bugs and makes the code harder to maintain. Could you please explain why these files are being excluded? If there are specific type errors that are difficult to resolve, it might be better to address them with # type: ignore comments and a TODO, rather than excluding the entire files.