nemo_automodel.components.speculative.eagle.remote.client

Training-side client for the remote EAGLE-3 target server.

RemoteEagle3TargetModel implements the Eagle3TargetBackend contract by delegating generate_batch to one or more remote target servers. It POSTs input_ids over HTTP and receives the supervision tensors either over NCCL (GPU-direct, body carries only metadata) or as a binary wire blob (fallback).

Multiple server URLs are dispatched round-robin so the prefetch pipeline in the training loop can keep several requests in flight (one per server) and overlap target inference with draft training.

Module Contents

Classes

Name	Description
`RemoteEagle3TargetModel`	EAGLE-3 target backend that delegates forward passes to remote servers.
`RemoteTargetPool`	Round-robin connections to the target servers of one remote training backend.
`TargetServerClient`	HTTP + NCCL connection to a single remote target server.
`_AsyncHandle`	Future-like wrapper that converts a worker-thread result into a batch.

Data

logger

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.

_embeddings

SimpleNamespace | None = None

_executor

ThreadPoolExecutor | None = None

_pool

num_remote_servers

int

supports_async

bool

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

class nemo_automodel.components.speculative.eagle.remote.client.RemoteTargetPool(
    urls: list[str],
    timeout: int,
    max_retries: int,
    enable_nccl: bool = True
)

Round-robin connections to the target servers of one remote training backend.

Owns what every remote target backend shares: the per-server clients, server selection, input encoding, weight fetching, and shutdown.

Parameters:

urls

list[str]

One http://host:port URL per target server.

timeout

int

Per-request HTTP timeout in seconds.

max_retries

int

Retries for idempotent requests on connection errors.

enable_nccl

boolDefaults to True

Whether the clients may use the NCCL data plane; see TargetServerClient.

_next

= itertools.cycle(self.clients)

clients

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

Disconnect from every server; a server shuts down when a client disconnects.

nemo_automodel.components.speculative.eagle.remote.client.RemoteTargetPool.encode_inputs(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor
) -> bytes

staticmethod

Encode one /generate request.

Parameters:

input_ids

torch.Tensor

Long tensor of shape [batch, sequence] containing target tokens.

attention_mask

torch.Tensor

Binary tensor of shape [batch, sequence], one for valid tokens.

loss_mask

torch.Tensor

Tensor of shape [batch, sequence] selecting supervised tokens.

Returns: bytes

Wire bytes holding CPU copies of the three tensors.

nemo_automodel.components.speculative.eagle.remote.client.RemoteTargetPool.fetch_weight(
    endpoint: str
) -> torch.Tensor

Fetch one full target weight, of shape [vocab, hidden], to CPU from the first server.

nemo_automodel.components.speculative.eagle.remote.client.RemoteTargetPool.next_client() -> nemo_automodel.components.speculative.eagle.remote.client.TargetServerClient

Return the server to send the next request to (round-robin).

class nemo_automodel.components.speculative.eagle.remote.client.TargetServerClient(
    url: str,
    timeout: int,
    max_retries: int,
    nccl_rank_offset: int = 0,
    enable_nccl: bool = True
)

HTTP + NCCL connection to a single remote target server.

enable_nccl=False keeps this client on the wire data plane even when NEMO_EAGLE_ENABLE_NCCL allows NCCL; the server keeps one NCCL group, so only one client per server can use it.

_generate_lock

= threading.Lock()

_nccl

NCCLTransport | None = None

_nccl_enabled

_nccl_lock

= threading.Lock()

_session

= requests.Session()

url

= url.rstrip('/')

nemo_automodel.components.speculative.eagle.remote.client.TargetServerClient._host() -> str

nemo_automodel.components.speculative.eagle.remote.client.TargetServerClient._init_nccl() -> bool

nemo_automodel.components.speculative.eagle.remote.client.TargetServerClient._nccl_port() -> int

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

nemo_automodel.components.speculative.eagle.remote.client.TargetServerClient.generate(
    payload: bytes
) -> dict[str, torch.Tensor | None]

POST /generate and return the supervision tensors (NCCL or wire).

/generate is the per-step hot path, so a transient timeout / connection reset here would otherwise abort a long remote-training run. The wire path is an idempotent HTTP round-trip, so it reuses request’s exponential-backoff retry. The NCCL path is deliberately a single attempt: the POST triggers a server-side NCCL send paired with the recv_tensors below, so a blind retry would issue a second send and desync the 2-process data-plane group (the client’s one recv vs the server’s two sends would hang). Recovering the NCCL path needs a transport resync (tear down + re-init, or fall back to wire) and is tracked separately.

Serialized on _generate_lock so this process never has two /generate requests in flight against the same server: the NCCL recv posted here must pair with this request’s send, and the server’s hook-based aux capture is not reentrant.

nemo_automodel.components.speculative.eagle.remote.client.TargetServerClient.request(
    endpoint: str,
    payload: bytes,
    content_type: str = 'application/octet-stream'
) -> bytes

POST payload to endpoint with exponential-backoff retry.

class nemo_automodel.components.speculative.eagle.remote.client._AsyncHandle(
    future,
    convert
)

Future-like wrapper that converts a worker-thread result into a batch.

nemo_automodel.components.speculative.eagle.remote.client._AsyncHandle.cancel() -> bool

nemo_automodel.components.speculative.eagle.remote.client._AsyncHandle.result(
    timeout: float | None = None
) -> nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch

nemo_automodel.components.speculative.eagle.remote.client.logger = logging.getLogger(__name__)