fix ppo value head load bugs

[lilei199908]

No description provided.

[Copilot]

Pull request overview

Adds a safeguard when initializing PPO critic models so the value head (output_layer) can be reinitialized after checkpoint load when the checkpoint metadata indicates it’s missing or shape-mismatched, avoiding incorrect/partial loads.

Changes:

• Add helpers to locate critic output_layer modules and detect whether they need reinitialization based on checkpoint tensor metadata.
• Reinitialize critic output_layer weights/bias after load_checkpoint() when required, and refresh optimizer param copies for fp16/bf16.

---

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

[Copilot]

_iter_critic_output_layers() yields chunk_id, but the returned chunk_id is never used by callers. This makes the helper harder to read than necessary. Either drop chunk_id from the yield, or include it in the warning/logging so it’s actionable when multiple pipeline chunks exist.

[Copilot]

The new checkpoint-metadata-based reinitialization path for the critic output layer is subtle and easy to regress (e.g., loading a checkpoint that doesn’t contain output_layer.* or has a different shape). It would be good to add a regression test that exercises loading a critic from a checkpoint missing/mismatching the value head and asserts training can proceed and the output layer parameters are reinitialized as intended.

Suggested change

	optimizer.reload_model_params()
	reload_model_params = getattr(optimizer, "reload_model_params", None)
	if callable(reload_model_params):
	reload_model_params()
	else:
	logger.warning(
	"Critic output layer was reinitialized after checkpoint load, but optimizer %s "
	"does not implement reload_model_params(); continuing without refreshing "
	"mixed-precision optimizer parameter copies.",
	type(optimizer).__name__,
	)

[Copilot]

_critic_output_layer_needs_reinit() does an O(N) scan over the entire checkpoint_metadata dict for each parameter via next(... for key, tensor_metadata in checkpoint_metadata.items() ...). For large distributed checkpoints this can add noticeable startup latency. Consider trying an exact-key lookup first (e.g., checkpoint_metadata.get(param_name)), and only falling back to a suffix scan if needed, or building a one-time suffix->metadata index to avoid repeated full scans.