Skip to content

fix(actor): use training_done event under Fast-LLM, not sample-counting heuristic #142

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
RaymondLi0 wants to merge 2 commits into
base: fast-llm
Choose a base branch
from
Open

fix(actor): use training_done event under Fast-LLM, not sample-counting heuristic #142

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7c2f75e
fix(actor): use training_done event under Fast-LLM, not sample-counti…
RaymondLi0 May 12, 2026
9ab0e53
review: trim verbose comment on is_trainer_finished
RaymondLi0 May 14, 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
20 changes: 17 additions & 3 deletions pipelinerl/actor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,11 @@ async def schedule_rollouts(
retry_max_delay_s = float(getattr(cfg.actor, "rollout_retry_max_delay_s", 30.0))

def is_trainer_finished() -> bool:
# Fast-LLM ignores `gradient_accumulation_passes` and overshoots `docs_per_step`
# by a few docs per step, so the sample-counting formula below fires several
# optimizer steps early. Use the explicit `training_finished` event instead.
if cfg.use_fast_llm:
return trainer_state.training_done
return (
trainer_state.samples_processed is not None
and trainer_state.samples_processed >= samples_target
Expand Down Expand Up @@ -609,9 +614,18 @@ def _run(self, dataset: list[tuple[str, dict]]):
# the user function must do next(...) to run each iteration
yield

final_steps = calculate_train_steps(self.cfg.finetune, self.cfg.finetune.interrupt_train_steps)
samples_target = final_steps * self.cfg.finetune.train_batch_size * self.cfg.finetune.gradient_accumulation_passes
if self.trainer_state.samples_processed is not None and self.trainer_state.samples_processed >= samples_target:
# Mirror `is_trainer_finished` (above): use the explicit training_done
# event under Fast-LLM; fall back to sample counting for HF/DeepSpeed.
if self.cfg.use_fast_llm:
trainer_finished = self.trainer_state.training_done
else:
final_steps = calculate_train_steps(self.cfg.finetune, self.cfg.finetune.interrupt_train_steps)
samples_target = final_steps * self.cfg.finetune.train_batch_size * self.cfg.finetune.gradient_accumulation_passes
trainer_finished = (
self.trainer_state.samples_processed is not None
and self.trainer_state.samples_processed >= samples_target
)
if trainer_finished:
logger.info("Trainer signalled completion; stopping actor loop")
break

Expand Down