`nemo_rl.models.generation.trtllm.config`#

Module Contents#

Classes#

`TrtllmSpecificArgs`
`TrtllmConfig`

API#

class nemo_rl.models.generation.trtllm.config.TrtllmSpecificArgs#

Bases: typing.TypedDict

tensor_parallel_size: int#: None

model_name: NotRequired[str]#: None

gpu_memory_utilization: NotRequired[float]#: None

max_model_len: int#: None

precision: str#: None

max_batch_size: int#: None

max_num_tokens: int#: None

expose_http_server: NotRequired[bool]#: None

async_engine: NotRequired[bool]#: None

moe_tensor_parallel_size: NotRequired[int]#: None

moe_expert_parallel_size: NotRequired[int]#: None

in_flight_weight_updates: NotRequired[bool]#: None

recompute_kv_cache_after_weight_updates: NotRequired[bool]#: None

default_chat_template_kwargs: NotRequired[dict[str, Any]]#: None

tool_parser: NotRequired[str]#: None

reasoning_parser: NotRequired[str]#: None

class nemo_rl.models.generation.trtllm.config.TrtllmConfig#

Bases: nemo_rl.models.generation.interfaces.GenerationConfig

trtllm_cfg: nemo_rl.models.generation.trtllm.config.TrtllmSpecificArgs#: None

trtllm_kwargs: NotRequired[dict[str, Any]]#: None