`nemo_rl.models.megatron.memory_saver`#

Module Contents#

`inference_model_alloc_region`	Allocation region to build the colocated inference model under.
`pause_inference_weights`	Back the colocated inference weights to CPU (no-op without torch_memory_saver).
`resume_inference_weights`	Restore the inference weights to their GPU addresses (no-op without torch_memory_saver).

nemo_rl.models.megatron.memory_saver._INFERENCE_MODEL_OFFLOAD_TAG#: ‘nemo_rl_megatron_inference_model’

nemo_rl.models.megatron.memory_saver.inference_model_alloc_region() → ContextManager[None][source]#

Allocation region to build the colocated inference model under.

Returns a CPU-backup-enabled torch_memory_saver region, or a null context.

nemo_rl.models.megatron.memory_saver.pause_inference_weights() → None[source]#: Back the colocated inference weights to CPU (no-op without torch_memory_saver).

nemo_rl.models.megatron.memory_saver.resume_inference_weights() → None[source]#: Restore the inference weights to their GPU addresses (no-op without torch_memory_saver).