nemo_rl.models.generation.vllm.vllm_backend

Contenu du module

Classes

API

class nemo_rl.models.generation.vllm.vllm_backend.VllmInternalWorkerExtension

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

Initialiser la communication collective.

report_device_id() -> str

prepare_refit_info(state_dict_info: typing.Optional[dict[str, typing.Any]] = None) -> None

Préparer les informations pour le réajustement.

DtensorPolicyWorker : inference colocalisé : state_dict_info est None inference non colocalisé : state_dict_info est un dictionnaire de {nom_du_tenseur : (forme, type_de_données)}

MegatronPolicyWorker : inference colocalisé : state_dict_info est un dictionnaire de {nom_du_tenseur : (forme, type_de_données, nombre_d’éléments)} inference non colocalisé : state_dict_info est un dictionnaire de {nom_du_tenseur : (forme, type_de_données)}

update_weights_from_global_ipc_handles(global_device_ipc_handles)

Mettre à jour les poids à partir des handles IPC globaux.

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

Returns : bool : True si les poids ont été mis à jour avec succès.

update_weights_from_local_ipc_handles(local_device_ipc_handles)

Mettre à jour les poids à partir des handles IPC locaux.

Args : local_device_ipc_handles (dict) : Handles IPC des paramètres pour le périphérique local.

Returns : bool : True si les poids ont été mis à jour avec succès.

update_weights_from_collective() -> bool

Mettre à jour les poids du modèle à partir de la communication collective.

start_gpu_profiling() -> None

Démarrer le profilage GPU.

stop_gpu_profiling() -> None

Arrêter le profilage GPU.