nemo_automodel.components.speculative.eagle.remote

Train-inference disaggregation for EAGLE-3 target serving.

Runs the frozen target model as a standalone inference server on separate GPU(s) while the draft model trains elsewhere. The training side talks to the server through RemoteEagle3TargetModel, which implements the Eagle3TargetBackend contract, so the EAGLE-3 recipe consumes a remote target exactly like the co-located HFEagle3TargetModel.

HTTP is the control plane (input_ids up, tensor metadata down).
NCCL is the data plane for the large supervision tensors (GPU-to-GPU, NVLink intra-node / RDMA inter-node), with a compact binary wire format fallback when NCCL is unavailable.

Submodules

Package Contents

Classes

Name	Description
`RemoteEagle3TargetModel`	EAGLE-3 target backend that delegates forward passes to remote servers.

API

class nemo_automodel.components.speculative.eagle.remote.client.RemoteEagle3TargetModel(
    urls: list[str],
    device: torch.device,
    timeout: int = 120,
    max_retries: int = 3
)

Bases: Eagle3TargetBackend

EAGLE-3 target backend that delegates forward passes to remote servers.

_clients

_embeddings

SimpleNamespace | None = None

_executor

ThreadPoolExecutor | None = None

_next

= itertools.cycle(range(len(self._clients)))

num_remote_servers

int

supports_async

bool

nemo_automodel.components.speculative.eagle.remote.client.RemoteEagle3TargetModel._build_payload(
    input_ids,
    attention_mask,
    loss_mask
) -> bytes

staticmethod

nemo_automodel.components.speculative.eagle.remote.client.RemoteEagle3TargetModel._to_batch(
    result: dict,
    attention_mask: torch.Tensor
) -> nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch

nemo_automodel.components.speculative.eagle.remote.client.RemoteEagle3TargetModel.close() -> None

nemo_automodel.components.speculative.eagle.remote.client.RemoteEagle3TargetModel.from_urls(
    urls: list[str],
    device,
    kwargs = {}
) -> 'RemoteEagle3TargetModel'

classmethod

nemo_automodel.components.speculative.eagle.remote.client.RemoteEagle3TargetModel.generate_batch(
    input_ids,
    attention_mask,
    loss_mask
) -> nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch

nemo_automodel.components.speculative.eagle.remote.client.RemoteEagle3TargetModel.generate_batch_async(
    input_ids,
    attention_mask,
    loss_mask
) -> nemo_automodel.components.speculative.eagle.remote.client._AsyncHandle

nemo_automodel.components.speculative.eagle.remote.client.RemoteEagle3TargetModel.get_input_embeddings()

Fetch the target input-embedding weight once and cache it.

Returns an object exposing .weight (the only attribute the draft’s copy_embeddings_from_target reads), matching the offline-cache path.

nemo_automodel.components.speculative.eagle.remote.client.RemoteEagle3TargetModel.model_info() -> dict

nemo_automodel.components.speculative.eagle.remote.client.RemoteEagle3TargetModel.set_vocab_mapping(
    selected_token_ids: torch.Tensor,
    selected_token_mask: torch.Tensor
) -> None