`RLLearner` resume skips unconsumed data when `num_iterations > 1`

[yixiaoer]

I found this while testing GRPO checkpoint/resume with num_iterations=2. The fast-forward logic is in the base RLLearner, so the problem is not specific to GRPOLearner. It is independent of #2288: the reproducer below uses data_shuffle_seed=None.

Expected Behavior

After a fresh-process resume, the learner should continue from the first input micro-batch that was not consumed before the checkpoint:

next_input_micro_batch(resumed) == next_input_micro_batch(uninterrupted)

Actual Behavior

On resume, RLLearner initializes its fast-forward target from actor_trainer.iter_steps:

• _last_iter_step is initialized from actor_trainer.iter_steps:

  tunix/tunix/rl/rl_learner.py

  Line 154 in caa444b

  self._last_iter_step = self.rl_engine.actor_trainer.iter_steps

• _prepare_data fast-forwards the input iterator until that count is reached:

  tunix/tunix/rl/rl_learner.py

  Lines 398 to 407 in caa444b

  	try:
  	while True:
  	while (
  	mode == rl_engine_lib.Mode.TRAIN
  	and self._iter_steps < self._last_iter_step
  	): # fast forward the iterator if loading from a previous checkpoint.
  	next(iterator)
  	self._iter_steps += 1
  	if self._iter_steps == self._last_iter_step:
  	logging.info("Fast forwarded %d micro-batches.", self._iter_steps)

iter_steps counts trainer loop iterations, not consumed input micro-batches:

• it is documented as the "# of times trainer has looped":

  tunix/tunix/sft/peft_trainer.py

  Line 406 in caa444b

  self._iter_steps = 0 # represent # of times trainer has looped

• it is incremented once per trainer loop:
  

  tunix/tunix/sft/peft_trainer.py

  Line 1000 in caa444b

  self._iter_steps += 1

• on restore, it is reconstructed as train_steps * gradient_accumulation_steps:

  tunix/tunix/sft/peft_trainer.py

  Line 421 in caa444b

  self._iter_steps = self._train_steps * self.config.get_with_default(

With num_iterations = μ, RLLearner passes batch_repeat=self._num_iterations() and enqueues each produced example μ times:

• repeated enqueue:
  

  tunix/tunix/rl/rl_learner.py

  Lines 347 to 354 in caa444b

  	for _ in range(repeats):
  	if self._data_shuffle_seed is not None:
  	shuffle_seed, self._data_shuffle_seed = jax.random.split(
  	self._data_shuffle_seed
  	)
  	shuffled_indices = jax.random.permutation(shuffle_seed, len(examples))
  	for i in shuffled_indices:
  	data_queue.put([examples[i]])

• batch_repeat=self._num_iterations():
  

  tunix/tunix/rl/rl_learner.py

  Line 732 in caa444b

  batch_repeat=self._num_iterations(),

In the non-packing configuration used by the reproducer, each consumed input micro-batch therefore produces μ trainer inputs, so the trainer advances iter_steps μ times per consumed input micro-batch.

On resume, _prepare_data uses that trainer-loop count as the number of input micro-batches to fast-forward. With num_iterations=1, these counts coincide in this configuration. With num_iterations>1, they do not.

Minimal reproduction

This test uses the real RLLearner.__init__ (which sets _last_iter_step from actor_trainer.iter_steps) and the real _prepare_data, including its resume fast-forward. It uses data_shuffle_seed=None, num_iterations=2, and one input micro-batch per step.

The mock only supplies the restored counter value used by the real trainer path. With num_iterations=2 and gradient_accumulation_steps=1, processing A results in train_steps=2, so the restored iter_steps is 2 * 1 = 2.

Save as test_resume.py:

from unittest import mock

import numpy as np
from tunix.rl import algorithm_config, rl_learner
from tunix.rl.queue import data_queue

NUM_ITERATIONS = 2
# What a trainer restored after training on A reports:
# iter_steps = train_steps * gradient_accumulation_steps = 2 * 1.
RESTORED_ITER_STEPS = 2


class Learner(rl_learner.RLLearner):
  _generate_and_compute_advantage = lambda self, x, mode: x
  _compute_trajectory_ids = lambda self, x, steps: [""] * len(x["prompts"])
  _num_iterations = lambda self: NUM_ITERATIONS
  _num_generations = lambda self: 1


def make_learner(iter_steps=0):
  engine = mock.MagicMock()
  engine.cluster_config.training_config.max_seq_token_per_tpu = None
  engine.actor_trainer.model = None
  engine.actor_trainer.iter_steps = iter_steps
  return Learner(engine, algorithm_config.AlgorithmConfig(),
                 reward_fns=lambda **kw: [])


def prepare(learner, iterator):
  q = data_queue.SimpleDataQueue(maxsize=0)
  learner._prepare_data(iterator, proceed_num_steps=1, sample_repeat=1,
                        batch_repeat=NUM_ITERATIONS,
                        service_target_batch_size=1, data_queue=q)
  return [str(x[0]["prompts"][0]) for x in iter(q.get, None)]


def test_resume_continues_from_next_micro_batch():
  batches = lambda: iter([{"prompts": np.array([x])} for x in "ABC"])

  continuous = make_learner()
  iterator = batches()
  assert prepare(continuous, iterator) == ["A", "A"]  # 2 trainer inputs
  expected = prepare(continuous, iterator)

  resumed = make_learner(iter_steps=RESTORED_ITER_STEPS)
  actual = prepare(resumed, batches())

  print("\nuninterrupted next:", expected)
  print("resumed next:      ", actual)
  assert actual == expected

Run:

JAX_PLATFORMS=cpu python -m pytest -q -s test_resume.py

Observed output:

uninterrupted next: ['B', 'B']
resumed next:       ['C', 'C']

E     AssertionError: assert ['C', 'C'] == ['B', 'B']

Before the checkpoint, only one input micro-batch (A) was consumed, although it produced two trainer iterations because num_iterations=2. The uninterrupted run therefore continues with B. On resume, _last_iter_step=2 makes _prepare_data call next(iterator) twice. That skips A and also B, which was never consumed, so training resumes from C.

Additional validation

I reproduced the same behavior through the real RLEngine + Trainer + Orbax path on CPU (GRPOLearner, num_iterations=2, return_logprobs=True, one identifiable prompt per input micro-batch, prompts recorded in the reward function):

uninterrupted input micro-batches : ['A', 'B', 'C']
before checkpoint                : ['A'] | train_steps 2
restored train_steps, iter_steps : 2 2
resumed processes                : ['C']

The real restored iter_steps matches the value used in the reproducer. With num_iterations=1 and everything else the same, the resumed run processes ['B', 'C'] as expected.

Environment

• CPU reproducer (JAX_PLATFORMS=cpu)
• originally observed on TPU during GRPO checkpoint/resume testing with num_iterations=2
• Tunix main @ caa444b73ffb1eb28b706f82c152fbe6ad7d5cd3, google-tunix 0.1.8
• JAX 0.11.0
• Flax 0.12.8
• Python 3.12

Checklist

• I have searched the existing issues for a similar bug report.
• I have provided all the required information in the "Environment" section.
• I have provided a minimal, reproducible example.

Would you like to help us fix it?
Yes, I'm happy to work on a fix. I'd like to confirm the expected resume behavior when a checkpoint is saved between num_iterations updates.