Skip to content

Remove LR scheduler for workers#1002

Open
jonahsamost wants to merge 3 commits intoNovaSky-AI:mainfrom
jonahsamost:jonah_2_1_scheduler
Open

Remove LR scheduler for workers#1002
jonahsamost wants to merge 3 commits intoNovaSky-AI:mainfrom
jonahsamost:jonah_2_1_scheduler

Conversation

@jonahsamost
Copy link

@jonahsamost jonahsamost commented Feb 1, 2026

Referencing #998

Remove internal learning rate scheduler logic from the skyRL train workers.

Since #978 added set_lr() support, Tinker can now control the LR externally. Since the workers's LRs could conflict with Tinker's LR, we just remove the ability for workers to set their own LRs

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Code Review

This pull request effectively removes the internal learning rate scheduler logic from the training workers, centralizing LR control under an external system like Tinker. The changes are consistent and thorough across FSDP and Megatron strategies, abstract base classes, and worker implementations. Key modifications include removing scheduler instantiation, step() calls, and state management from checkpoints. The code is now cleaner and the responsibility for learning rate scheduling is clearly externalized. The changes look solid.

Comment on lines 21 to 24
from skyrl_train.distributed.megatron.optimizer import (
init_megatron_optim_config,
get_megatron_optimizer,
get_megatron_optimizer_param_scheduler,
)
Copy link
Contributor

Choose a reason for hiding this comment

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

medium

With the removal of get_megatron_optimizer_param_scheduler from the imports, the function itself in skyrl_train/skyrl_train/distributed/megatron/optimizer.py appears to be dead code. Consider removing it in a follow-up to improve code maintainability.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant