nemo_rl.models.generation.megatron.utils#

Module Contents#

Functions#

sample_vision_tensors

Return one sample’s vision tensors from RL PackedTensors.

build_prompt_and_multimodal_data

Build one pre-expanded token prompt and optional MCore media payload.

build_image_preprocessing_config

Translate an HF image processor to an MCore config.

build_video_preprocessing_config

Build video preprocessing when explicitly enabled by generation config.

resolve_torch_dtype

Convert a value to torch.dtype.

log_gpu_memory

Print a one-line GPU-memory summary for the calling rank.

API#

nemo_rl.models.generation.megatron.utils.sample_vision_tensors(data, index: int)#

Return one sample’s vision tensors from RL PackedTensors.

nemo_rl.models.generation.megatron.utils.build_prompt_and_multimodal_data(
data,
index: int,
*,
supports_modality: Callable[[str], bool],
sample_tensors: Callable = sample_vision_tensors,
)#

Build one pre-expanded token prompt and optional MCore media payload.

nemo_rl.models.generation.megatron.utils.build_image_preprocessing_config(
image_processor: Any,
*,
dynamic_resolution: bool | None = None,
vision_model_type: str | None = None,
dynamic_resolution_model_length: int | None = None,
dynamic_resolution_rounding_mode: str | None = None,
dynamic_resolution_resize_mode: str | None = None,
) → megatron.core.inference.config.ImageProcessingConfig#

Translate an HF image processor to an MCore config.

Parameters:

  • image_processor – HF image processor to read patch/normalization fields from.

  • dynamic_resolution – Override for ImageProcessingConfig.dynamic_resolution. None leaves MCore’s own default in place.

  • vision_model_type – Override for ImageProcessingConfig.vision_model_type. None leaves MCore’s own default in place.

  • dynamic_resolution_model_length – Optional override for the processor’s max_model_len.

  • dynamic_resolution_rounding_mode – Optional patch-grid rounding contract.

  • dynamic_resolution_resize_mode – Optional image-resize contract.

nemo_rl.models.generation.megatron.utils.build_video_preprocessing_config(
image_config: megatron.core.inference.config.ImageProcessingConfig | None,
generation_config: dict[str, Any],
*,
frame_manifest_magic: bytes,
) → megatron.core.inference.config.VideoProcessingConfig | None#

Build video preprocessing when explicitly enabled by generation config.

nemo_rl.models.generation.megatron.utils.resolve_torch_dtype(val)#

Convert a value to torch.dtype.

nemo_rl.models.generation.megatron.utils.log_gpu_memory(tag: str) → None#

Print a one-line GPU-memory summary for the calling rank.