`nemo_rl.models.generation`#

Subpackages#

Submodules#

Package Contents#

Functions#

`maybe_configure_engine_reaping_env`	Let the raylet reap a dead generation worker’s EngineCore, on runs the driver owns.
`resolve_generation_class`	Map `generation_config` to its GenerationInterface class.
`draft_full_refit_enabled`	Whether the trainer streams the drafter’s FULL weight set on refit.
`configure_generation_config`	Apply specific configurations to generation config.

Data#

TokenizerType

API#

nemo_rl.models.generation.TokenizerType#: None

nemo_rl.models.generation.maybe_configure_engine_reaping_env(enabled: bool) → None#

Let the raylet reap a dead generation worker’s EngineCore, on runs the driver owns.

A vLLM generation worker spawns its EngineCore as a plain multiprocessing child – context.Process(target=EngineCoreProc.run_engine_core), no setsid, not a daemon – and that child owns the CUDA context and the KV cache. Every cleanup path that exists runs INSIDE the dying process, so SIGKILL defeats all of them: vLLM’s weakref.finalize never fires, a non-daemon multiprocessing child is designed to outlive its parent, and Ray’s default RAY_kill_child_processes_on_worker_exit is the worker’s own exit handler. The orphan then holds its GPU for the life of the job: measured on 4xGB200, one SIGKILLed shard left 114.95 GiB of a 184.31 GiB device still held, and it did not come back across five restart attempts – the replacement could not fit, so re-admission failed while the survivors carried on. Because the EngineCore never calls setsid it stays in the worker’s process group, which is exactly what per-worker process-group cleanup reaches.

The raylet is the only reader, and this only reaches the raylet on runs where the driver starts it. Ray reads process_group_cleanup_enabled in the raylet, at the raylet’s own startup. Call this before init_ray(), like

Func:: ~nemo_rl.data_plane.factory.maybe_configure_data_plane_env, and on a local/single-node run ray.init forks a raylet that inherits this environment, so it lands. On a Slurm or k8s cluster it does NOT: ray start has already run on every node before the driver process exists, and init_ray() then attaches to it with address="auto". Nothing the driver sets afterwards can change a raylet that is already up. Those deployments set the flag where the raylet is actually launched – see the RAY_process_group_cleanup_enabled export in ray.sub – which means it is unconditional there, not scoped to fleet health.

So the gate below buys scoping on exactly one deployment shape. It is still worth having: a raylet-wide behaviour change should not follow from import nemo_rl, and a plain SFT, DPO or distillation run has no EngineCore to reap. But it is not a property of the flag, it is a property of who started the raylet.

setdefault, not assignment: an operator debugging a wedged engine may want the corpse, and an explicit RAY_process_group_cleanup_enabled=0 stays honoured.

Parameters:: enabled – Whether this run can lose a generation worker and keep going. A bool rather than the fleet-health config, so any entrypoint that builds vLLM generation can call it – the leak is not specific to fleet health, it just takes a run that survives the loss to leave the orphan behind. Note this is fleet_health.enabled, not restart_dead_shards: any shard loss leaks that GPU, restarting is only what makes the loss survivable.

nemo_rl.models.generation.resolve_generation_class( generation_config: nemo_rl.models.generation.interfaces.GenerationConfig, ) → type[nemo_rl.models.generation.interfaces.GenerationInterface]#: Map generation_config to its GenerationInterface class.

nemo_rl.models.generation.draft_full_refit_enabled(policy_cfg: dict) → bool#

Whether the trainer streams the drafter’s FULL weight set on refit.

True for Automodel draft co-training (dspark/dflash/eagle3), which exports the drafter’s entire state_dict under the draft. prefix; the megatron eagle3 path streams a partial set instead. Feeds the draft_full_refit argument of :func:configure_generation_config.

automodel_cfg.enabled alone identifies the Automodel path.

nemo_rl.models.generation.configure_generation_config( config: nemo_rl.models.generation.interfaces.GenerationConfig, tokenizer: nemo_rl.models.generation.TokenizerType, is_eval: bool = False, has_refit_draft_weights: bool = False, trains_mtp: bool = False, draft_full_refit: bool = False, ) → nemo_rl.models.generation.interfaces.GenerationConfig#: Apply specific configurations to generation config.

nemo_rl.models.generation#