refactor/ppo

[lilei199908]

future TODO:

1. maybe not need foward_only for values.
2. curret only surpot same parapall between actor and critic, maybe need more flexible.
3. mayne can considerate muti ciritc
4. more test

curret reward compare verl
up slime， down verl

[Copilot]

Pull request overview

Refactors PPO actor/critic training to decouple critic configuration from the actor and to move actor↔critic value sharing out of NCCL process-group broadcast and into Ray-mediated data passing (with DP/CP shard identity).

Changes:

• Adds critic YAML override parsing (--critic-config-path) and a --custom-advantage-fn hook.
• Reworks Ray training actors/groups to support passing per-worker external_data (e.g., critic values) into actor training.
• Removes the actor↔critic NCCL “connect” path and related broadcast sync, replacing it with CPU tensor transfer helpers and shard-identity APIs.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
slime/utils/misc.py	Adds helper to map critic value shards to actor workers by (DP, CP).
slime/utils/arguments.py	Introduces critic YAML parsing + new CLI flags; removes old critic-only CLI flags; adjusts critic GPU assumptions.
slime/ray/train_actor.py	Extends train() signature and exposes per-worker parallel value identity.
slime/ray/rollout.py	Updates GPU offset/size computations (now excluding critic).
slime/ray/placement_group.py	Splits critic into its own placement group and wires critic arg overrides into model creation.
slime/ray/actor_group.py	Adds external_data plumbing to async_train() and exposes group shard identities.
slime/backends/megatron_utils/loss.py	Adds custom_advantage_fn hook and persists computed KL into rollout data.
slime/backends/megatron_utils/data.py	Replaces actor↔critic broadcast sync with CPU↔GPU tensor transfer helpers.
slime/backends/megatron_utils/arguments.py	Removes critic-train-only world-size special-casing (world size now actor-based).
slime/backends/megatron_utils/actor.py	Removes actor↔critic NCCL connect; critic returns CPU values; actor optionally consumes external CPU values.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

critic_values_by_actor_worker will raise KeyError for actor workers that don't have a corresponding entry in shard_to_values (e.g., non-last PP stages, since critic train_critic() returns {} for those workers). Consider returning {} for workers with no values (or only generating entries for last-PP-stage workers), and validate the lengths/world-sizes so missing shards fail fast instead of being silently truncated by zip(..., strict=False).

Suggested change

	shard_to_values = {
	(dp, cp): value["values"]
	for (dp, cp, _, _), value in zip(critic_model.get_parallel_value_info(), per_worker_values, strict=False)
	if "values" in value
	}
	return [{"values": shard_to_values[(dp, cp)]} for dp, cp, _, _ in actor_model.get_parallel_value_info()]
	critic_parallel_info = list(critic_model.get_parallel_value_info())
	actor_parallel_info = list(actor_model.get_parallel_value_info())
	if len(critic_parallel_info) != len(per_worker_values):
	raise ValueError(
	"critic parallel info and per_worker_values must have the same length: "
	f"{len(critic_parallel_info)} != {len(per_worker_values)}"
	)
	critic_shards = {(dp, cp) for dp, cp, _, _ in critic_parallel_info}
	actor_shards = {(dp, cp) for dp, cp, _, _ in actor_parallel_info}
	missing_actor_shards = actor_shards - critic_shards
	if missing_actor_shards:
	raise ValueError(
	"actor workers require (dp, cp) shards that do not exist in critic parallel info: "
	f"{sorted(missing_actor_shards)}"
	)
	shard_to_values = {
	(dp, cp): value["values"]
	for (dp, cp, _, _), value in zip(critic_parallel_info, per_worker_values)
	if "values" in value
	}
	return [
	{"values": shard_to_values[(dp, cp)]} if (dp, cp) in shard_to_values else {}
	for dp, cp, _, _ in actor_parallel_info
	]

[Copilot]

Creating a separate placement group for the critic removes gang-scheduling with the actor+rollout group. Because the actor+rollout PG is created (and pg.ready() awaited) before the critic PG, it's possible to reserve resources in a way that prevents the critic PG from ever being scheduled (resource fragmentation/deadlock when the cluster is near capacity). Consider allocating a single placement group for actor+critic+rollout again, or otherwise ensuring atomic/gang allocation (e.g., one PG with bundle offsets for each role, or a strategy that avoids partial allocation).

[Copilot]

parse_critic_args uses assert for user config validation. Asserts can be disabled with python -O, turning this into undefined behavior later. Prefer raising a ValueError (or RuntimeError) with the same message so validation is always enforced.

Suggested change

	assert isinstance(critic_entries, list) and len(critic_entries) == 1, (
	"critic config must contain exactly one entry under 'critic', e.g. "
	"critic: [{name: default, overrides: {...}}]"
	)
	if not (isinstance(critic_entries, list) and len(critic_entries) == 1):
	raise ValueError(
	"critic config must contain exactly one entry under 'critic', e.g. "
	"critic: [{name: default, overrides: {...}}]"
	)

[Copilot]

slime_validate_args now forces critic_num_nodes/gpus_per_node to always match the actor, but some downstream logic still treats critic GPUs as part of the rollout/colocate GPU budget. Double-check that rollout_num_gpus and placement-group sizing remain consistent under --colocate with --advantage-estimator ppo; otherwise rollout server sizing/validation can diverge from the actual rollout placement group.

Suggested change

	actor_num_gpus = args.actor_num_gpus_per_node * args.actor_num_nodes
	if args.colocate and args.use_critic:
	# Under colocated PPO, actor and critic share the same placement-group GPUs.
	# Keep rollout sizing aligned with the colocated actor/critic world size instead
	# of implicitly budgeting critic GPUs a second time.
	if not args.rollout_num_gpus:
	args.rollout_num_gpus = actor_num_gpus
	else:
	assert args.rollout_num_gpus == actor_num_gpus, (
	"When colocate is enabled with advantage_estimator=ppo, "
	"rollout_num_gpus must match the colocated actor/critic GPU count "
	f"({actor_num_gpus}), got {args.rollout_num_gpus}."
	)

[Copilot]

compute_advantages_and_returns assumes a custom advantage function will always populate rollout_data['advantages'] and rollout_data['returns'], but if it doesn't, the code will fail later with a KeyError/unclear error. Consider validating after custom_adv_fn(...) returns and raising a clear ValueError if required keys are missing or have the wrong type/shape.

Suggested change

	advantages, returns = rollout_data["advantages"], rollout_data["returns"]
	missing_keys = [
	key for key in ("advantages", "returns") if key not in rollout_data
	]
	if missing_keys:
	raise ValueError(
	f"Custom advantage function {args.custom_advantage_fn!r} must populate "
	f"rollout_data with keys 'advantages' and 'returns'; missing keys: "
	f"{missing_keys}."
	)
	advantages = rollout_data["advantages"]
	returns = rollout_data["returns"]
	if not isinstance(advantages, list) or not isinstance(returns, list):
	raise ValueError(
	f"Custom advantage function {args.custom_advantage_fn!r} must populate "
	f"rollout_data['advantages'] and rollout_data['returns'] as lists of "
	f"torch.Tensor objects, got {type(advantages).__name__} and "
	f"{type(returns).__name__}."
	)
	expected_num_samples = len(kl)
	if len(advantages) != expected_num_samples or len(returns) != expected_num_samples:
	raise ValueError(
	f"Custom advantage function {args.custom_advantage_fn!r} returned "
	f"{len(advantages)} advantages and {len(returns)} returns, expected "
	f"{expected_num_samples} of each."
	)
	for i, (advantage, ret, k) in enumerate(
	zip(advantages, returns, kl, strict=False)
	):
	if not isinstance(advantage, torch.Tensor) or not isinstance(ret, torch.Tensor):
	raise ValueError(
	f"Custom advantage function {args.custom_advantage_fn!r} must return "
	f"lists of torch.Tensor objects; sample {i} has types "
	f"{type(advantage).__name__} and {type(ret).__name__}."
	)
	if advantage.shape != k.shape or ret.shape != k.shape:
	raise ValueError(
	f"Custom advantage function {args.custom_advantage_fn!r} returned "
	f"incompatible tensor shapes for sample {i}: advantages shape "
	f"{tuple(advantage.shape)}, returns shape {tuple(ret.shape)}, "
	f"expected {tuple(k.shape)}."
	)

[Copilot]

This change removes --critic-train-only (and related critic-* CLI args), but train.py and train_async.py still read args.critic_train_only in multiple places. After this PR, those entrypoints will raise AttributeError unless they are updated (or the flag is kept with a deprecated alias).

[Copilot]

async_train() now supports external_data, but there are currently no call sites passing critic-produced values into the actor. With the removal of sync_actor_critic_data, PPO actor training will compute advantages without rollout_data['values'] and will fail in compute_advantages_and_returns() (PPO branch requires values). Wire the training loop to (1) ray.get() critic results, (2) map them to actor workers (DP+CP), and (3) pass per-worker external_data into actor_model.async_train(...).

[Copilot]

ignored_keys = {"num_nodes", "num_gpus_per_node"} doesn't match the actual CLI attribute names in this repo (e.g., actor_num_nodes, actor_num_gpus_per_node, critic_num_nodes, ...). As written, a critic YAML could override actor_num_nodes/actor_num_gpus_per_node inside critic_args, leading to inconsistent DP calculations and confusing behavior. Consider explicitly ignoring/forbidding overrides of the GPU allocation keys that should always follow the actor (and maybe validating that the YAML doesn't contain them).

Suggested change

	ignored_keys = {"num_nodes", "num_gpus_per_node"}
	# Apply overrides from the YAML config.
	# Unspecified keys inherit from actor_args via deepcopy.
	for key, value in critic_config.items():
	if key in ignored_keys:
	logger.info(f"Ignoring critic config key '{key}'; critic GPU allocation always follows actor.")
	continue
	forbidden_gpu_override_keys = {
	# Legacy generic names.
	"num_nodes",
	"num_gpus_per_node",
	# Actual CLI attribute names used by this repo.
	"actor_num_nodes",
	"actor_num_gpus_per_node",
	"critic_num_nodes",
	"critic_num_gpus_per_node",
	}
	invalid_gpu_override_keys = sorted(set(critic_config).intersection(forbidden_gpu_override_keys))
	if invalid_gpu_override_keys:
	raise ValueError(
	"Critic config cannot override GPU allocation keys that must follow the actor: "
	+ ", ".join(invalid_gpu_override_keys)
	)
	# Apply overrides from the YAML config.
	# Unspecified keys inherit from actor_args via deepcopy.
	for key, value in critic_config.items():

[Copilot]

_compute_megatron_num_gpus no longer includes critic GPUs, but its docstring still says "actor + critic". Update the docstring (and any dependent comments) to reflect the new placement-group layout so future changes don’t reintroduce incorrect offsets/sizing.

[zhuzilin]

Suggested change

	result["critic"] = None
	result["critic"] = result["actor"] if args.use_critic else None

[zhuzilin]

please follow the naming convention, use --custom-xxx-function-path

[zhuzilin]

why do we need to turn off normalize adv here?

[zhuzilin]

if we only supports setting critic and actor as the same parallel config, maybe we can simply pass the data from critic to actor and leave this change to future PRs.

[zhuzilin]

if we don't need critic_values_by_actor_worker, can we pass the value_refs to actor instread of materialize it in the control process?