nemo_rl.models.generation.vllm.vllm_worker_async

Contenu du module

Classes

API

class nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker

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

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

post_init_async()

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

generate_async(
    data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
    greedy: bool = False
) -> typing.AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]

Générer un lot de données à l’aide du moteur AsyncLLM de vLLM, en renvoyant les résultats dès qu’ils sont prêts.

Args : data : BatchedDataDict avec input_ids et input_lengths greedy : Indique s’il faut utiliser le décodage glouton au lieu de l’échantillonnage

Renvoie : Tuple de (indice_original, BatchedDataDict conforme à GenerationOutputSpec pour la séquence unique)

generate_text_async(
    data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
    greedy: bool = False
) -> typing.AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]

Générer de façon asynchrone des réponses textuelles, en renvoyant les résultats dès qu’ils sont prêts.

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

Renvoie : Tuple de (indice_original, BatchedDataDict contenant une seule réponse textuelle)

report_device_id_async() -> list[str]

Version asynchrone de report_device_id.

prepare_refit_info_async(state_dict_info: dict[str, typing.Any]) -> None

Version asynchrone de prepare_refit_info.

update_weights_from_ipc_handles_async(ipc_handles: dict[str, typing.Any]) -> bool

Version asynchrone de update_weights_from_ipc_handles.

Args : ipc_handles (dict) : Dictionnaire mappant les UUID de périphériques (str) aux handles IPC de paramètres.

Renvoie : bool : True si les poids ont été mis à jour avec succès, False sinon.

update_weights_from_collective_async() -> bool

Version asynchrone de update_weights_from_collective.

reset_prefix_cache_async()

Version asynchrone de reset_prefix_cache.

sleep_async()

Version asynchrone de sleep.

wake_up_async(**kwargs)

Version asynchrone de wake_up.

shutdown() -> bool

Nettoyer les ressources vLLM.