Checking main's conflicts with the Tmax branch

[shatu]

No description provided.

[gemini-code-assist]

Code Review

This pull request integrates OLMo-core's training infrastructure into the repository, adding a native SFT training script (olmo_core_finetune.py) and a dataset conversion utility (numpy_dataset_conversion.py) to prepare SFT datasets in numpy mmap format. It also refactors the DPO and GRPO trainers to align with OLMo-core utilities, introduces a unified ρ correction (train/infer ratio correction) in GRPO, updates evaluation scripts, and bumps dependencies. The review feedback identifies several key issues: guarding Ray actor calls in perform_weight_sync with is_rank_0 to prevent redundant RPC overhead, correcting a 0-based step indexing off-by-one error in EvalCallback, defensively handling numeric types in parse_dataset_mixer_list to avoid TypeError, and guarding logprob indexing in process_group to prevent potential IndexError crashes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In distributed training, perform_weight_sync is called on all ranks. However, Ray actor calls (such as pausing the actor manager and waking up the vLLM engines) are external, non-collective operations that should only be coordinated by rank 0. Having all ranks concurrently invoke these remote methods leads to massive redundant RPC overhead and potential bottlenecks/deadlocks. Guarding these calls with is_rank_0 ensures that only the main process interacts with the Ray actors.

Suggested change

	start = time.perf_counter()
	ray.get(actor_manager.set_should_stop.remote(True))
	try:
	results, actor_sync_times = utils.ray_get_with_progress(
	broadcast_refs, desc="Broadcasting weights to vLLM engines", enable=progress
	)
	if not inflight_updates:
	utils.ray_get_with_progress(
	itertools.chain.from_iterable(results), desc="Waiting for vLLM engine update RPCs", enable=progress
	)
	utils.ray_get_with_progress(
	[e.wake_up.remote() for e in vllm_engines], desc="Waking up vLLM engines", enable=progress
	)
	finally:
	ray.get(actor_manager.set_should_stop.remote(False))
	start = time.perf_counter()
	is_rank_0 = not dist.is_initialized() or dist.get_rank() == 0
	if not is_rank_0:
	return {"time/weight_sync": 0.0}, []
	ray.get(actor_manager.set_should_stop.remote(True))
	try:
	results, actor_sync_times = utils.ray_get_with_progress(
	broadcast_refs, desc="Broadcasting weights to vLLM engines", enable=progress
	)
	if not inflight_updates:
	utils.ray_get_with_progress(
	itertools.chain.from_iterable(results), desc="Waiting for vLLM engine update RPCs", enable=progress
	)
	utils.ray_get_with_progress(
	[e.wake_up.remote() for e in vllm_engines], desc="Waking up vLLM engines", enable=progress
	)
	finally:
	ray.get(actor_manager.set_should_stop.remote(False))

[gemini-code-assist]

In olmo_core, the trainer's global_step is 0-based, meaning the very first step of training is step 0. Checking self.trainer.global_step == 1 for eval_on_step_0 will cause the initial evaluation to run at step 1 instead of step 0 (or not run at all if step 0 is skipped). Updating the conditions to check for global_step == 0 and global_step > 0 ensures that the initial evaluation runs correctly at the start of training.

Suggested change

	if not (
	(self.args.eval_on_step_0 and self.trainer.global_step == 1)
	or (self.trainer.global_step % self.args.local_eval_every == 0 and self.trainer.global_step > 1)
	):
	if not (
	(self.args.eval_on_step_0 and self.trainer.global_step == 0)
	or (self.trainer.global_step % self.args.local_eval_every == 0 and self.trainer.global_step > 0)
	):

[gemini-code-assist]

If mixer_list is passed programmatically and already contains numeric types (such as float or int), performing "." in mixer_list[i + 1] will raise a TypeError because numbers are not iterable. Coercing or checking the type defensively prevents runtime crashes when the list is not entirely composed of strings.

Suggested change

	value = float(mixer_list[i + 1]) if "." in mixer_list[i + 1] else int(mixer_list[i + 1])
	raw_val = mixer_list[i + 1]
	if isinstance(raw_val, (int, float)):
	value = raw_val
	else:
	value = float(raw_val) if "." in raw_val else int(raw_val)

[gemini-code-assist]

If logprobs are not requested or not populated, result.logprobs may be empty, causing result.logprobs[i] to raise an IndexError. Adding a guard to check if i < len(result.logprobs) prevents potential runtime crashes when logprobs are disabled.

Suggested change

	for i in range(len(result.finish_reasons)):
	if result.finish_reasons[i] == "stop" and len(result.responses[i]) == 0:
	result.responses[i].append(tokenizer.eos_token_id)
	result.masks[i].append(1)
	result.logprobs[i].append(float("nan"))
	for i in range(len(result.finish_reasons)):
	if result.finish_reasons[i] == "stop" and len(result.responses[i]) == 0:
	result.responses[i].append(tokenizer.eos_token_id)
	result.masks[i].append(1)
	if i < len(result.logprobs):
	result.logprobs[i].append(float("nan"))