nemo_rl.models.automodel.utils

Module Contents

Functions

resolve_model_class	Resolve the model class for a model type.
release_moe_dispatcher_state	Drop the per-dispatch tensors cached on the model’s MoE dispatch managers.

Data

AUTOMODEL_FACTORY
_MOE_DISPATCH_STATE_ATTRS

API

nemo_rl.models.automodel.utils.AUTOMODEL_FACTORY: Dict[str, Any]

None

nemo_rl.models.automodel.utils.resolve_model_class(model_name: str) → Any

Resolve the model class for a model type.

Parameters:

model_name – Model type to resolve.

nemo_rl.models.automodel.utils._MOE_DISPATCH_STATE_ATTRS

(‘token_probs’, ‘token_indices’, ‘routing_map’, ‘handle’, ‘tokens_per_expert’, ‘num_permuted_tokens’…

nemo_rl.models.automodel.utils.release_moe_dispatcher_state(model: torch.nn.Module) → int

Drop the per-dispatch tensors cached on the model’s MoE dispatch managers.

The managers keep the last dispatch’s routing tensors as attributes; combine clears only handle (and HybridEP’s num_permuted_tokens). After a backward with activation checkpointing, a manager’s last dispatch is a recompute, so token_probs and dispatched_probs carry a grad_fn that keeps that recompute graph and its saved activations alive. With the default manager shared by all layers that is the first MoE layer (about 1 GB for Qwen3.5-35B-A3B at 9216 tokens per micro-batch); with per-layer managers every layer keeps one. The live blocks pin cached allocator segments, so torch.cuda.empty_cache() cannot return them before colocated generation wakes up.

Parameters:

model – Model whose token_dispatcher attributes are scanned. A manager shared across layers is cleared once.

Returns:

The number of attributes set to None.