nemo_rl.models.generation.vllm.vllm_worker

Module Contents

Classes

API

class nemo_rl.models.generation.vllm.vllm_worker.BaseVllmGenerationWorker(config: nemo_rl.models.generation.vllm.config.VllmConfig, bundle_indices: typing.Optional[list[int]] = None, fraction_of_gpus: float = 1.0, seed: typing.Optional[int] = None)

__repr__() -> str

Personnalise le préfixe de l’acteur dans les journaux de Ray.

Cela facilite l’identification du worker qui produit des messages de journal spécifiques.

configure_worker(
    num_gpus: int | float,
    bundle_indices: typing.Optional[tuple[int, list[int]]] = None
) -> tuple[dict[str, typing.Any], dict[str, str], dict[str, typing.Any]]

Fournit la configuration complète du worker pour le parallélisme tensoriel et par pipeline de vLLM.

Cette méthode configure le worker en fonction de son rôle dans le parallélisme tensoriel et par pipeline, qui est déterminé directement à partir du paramètre bundle_indices.

Args: num_gpus : Allocation GPU originale pour ce worker basée sur le groupe de placement bundle_indices : Tuple de (node_idx, local_bundle_indices) pour le parallélisme (le cas échéant)

Returns: tuple avec la configuration complète du worker :

• ‘resources’ : Allocation de ressources (par exemple, num_gpus)
• ‘env_vars’ : Variables d’environnement pour ce worker
• ‘init_kwargs’ : Paramètres à passer à init du worker

llm()

is_alive()

Vérifier si le worker est actif.

_merge_stop_strings(batch_stop_strings)

_build_sampling_params(
    *,
    greedy: bool,
    stop_strings,
    max_new_tokens: typing.Optional[int] = None
)

start_gpu_profiling() -> None

Démarrer le profilage GPU.

stop_gpu_profiling() -> None

Arrêter le profilage GPU.

class nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker(config: nemo_rl.models.generation.vllm.config.VllmConfig, bundle_indices: typing.Optional[list[int]] = None, fraction_of_gpus: float = 1.0, seed: typing.Optional[int] = None)

Bases : nemo_rl.models.generation.vllm.vllm_worker.BaseVllmGenerationWorker

_create_engine(llm_kwargs: dict[str, typing.Any]) -> None

post_init()

init_collective(
    rank_prefix: int,
    ip: str,
    port: int,
    world_size: int
) -> None

generate(
    data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
    greedy: bool = False
) -> nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]

Générer un lot de données à l’aide de la génération vLLM.

Args: data: BatchedDataDict contenant les tenseurs input_ids et input_lengths greedy: Indique s’il faut utiliser le décodage gourmand au lieu de l’échantillonnage

Returns: BatchedDataDict conforme à GenerationOutputSpec :

• output_ids : ID de tokens d’entrée + générés avec un remplissage approprié
• logprobs : Logarithmes des probabilités des tokens
• generation_lengths : Longueurs de chaque réponse
• unpadded_sequence_lengths : Longueurs de chaque séquence d’entrée + générée

generate_text(
    data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
    greedy: bool = False
) -> nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]

Générer des réponses textuelles à l’aide de la génération vLLM.

Args: data: BatchedDataDict contenant des invites avec des chaînes de texte greedy: Indique s’il faut utiliser le décodage gourmand au lieu de l’échantillonnage

Returns: BatchedDataDict contenant :

• texts : Liste de réponses textuelles générées