schedule free adamw jax update with switching params to x for validation, y for training

[wyfEmma]

New Submission

Submission Information

Please fill out the following information about your submission within the quotation marks.

submission_name: "Schedule-Free AdamW fix"  # As it will appear on the leaderboard
submission_folder: "schedule_free"  # Name of folder within `submissions/external_tuning/` or `submissions/self_tuning/` (lowercase, no spaces)
authors: "Yifan (Emma)"  # List authors separated by commas
affiliations: "Google"  # List all affiliations of the authors, separated by commas
version: "1.0"  # Optional version number of your submission
ruleset: "self-tuning"  # Either "external" or "self-tuning"
framework: "JAX"  # Either "PyTorch" or "JAX"
description: "schedule free adamw jax update with switching params to x for validation, y for training to match the algorithm, this fix is based off of init-22's pr"  # A short, high-level description of the algorithm

Evidence for the Submission's Performance

If possible provide some evidence of your submission's performance. E.g. a link to a paper/pre-print, training logs, screenshots, etc. The working group will prioritize evaluating submissions with more convincing evidence.

Comments

Feel free to add any comments, descriptions, or questions here.

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[priyakasimbeg]

Where is this implemented?

[wyfEmma]

this is implemented in the optax library for schedule free algorithm,
def schedule_free_eval_params(state: ScheduleFreeState, params: base.Params):
"""Params for evaluation of :func:optax.contrib.schedule_free."""
return jax.tree_util.tree_map(
lambda yi, zi: (yi - (1.0 - state.b1) * zi) / state.b1, params, state.z
)

[priyakasimbeg]

should we set this back to False after updating?

[priyakasimbeg]

nvm, got it now

[wyfEmma]

move 190-219 out of this function