Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsSpeculativeDsparknemo_automodel.components.speculative.dspark.remote

nemo_automodel.components.speculative.dspark.remote

View as Markdown

Train-inference disaggregation for DSpark: a remote target server and its client.

The frozen DSpark target runs in its own process, on other GPUs or nodes, and serves the two hidden-state tensors the draft trains on. The transport is the EAGLE-3 one, unchanged (nemo_automodel.components.speculative.eagle.remote): HTTP is the control plane and NCCL or the binary wire format is the data plane. Only the payload differs. The server captures target_hidden_states and target_last_hidden_states with HFDSparkTargetModel, the same code the colocated path runs, and both data planes transfer them losslessly, so a remote batch equals the colocated one for the same target forward.

Module Contents

Classes

NameDescription
DSparkTargetInfoTarget properties the draft is built against.
DSparkTargetModelServerRequest-handling logic for a remote DSpark target.
RemoteDSparkTargetModelDSpark target backend that runs the frozen target on remote servers.

API

class nemo_automodel.components.speculative.dspark.remote.DSparkTargetInfo(
    num_hidden_layers: int,
    hidden_size: int,
    vocab_size: int,
    target_layer_ids: tuple[int, ...]
)
Dataclass

Target properties the draft is built against.

The server reports them and the training client compares them with the values it derived from its own config, so a server running a different target, depth, or capture layers fails at setup instead of training on mismatched features.

hidden_size
int
num_hidden_layers
int
target_layer_ids
tuple[int, ...]
vocab_size
int
nemo_automodel.components.speculative.dspark.remote.DSparkTargetInfo.from_text_config(
    text_config: transformers.PretrainedConfig,
    target_layer_ids: collections.abc.Sequence[int]
) -> nemo_automodel.components.speculative.dspark.remote.DSparkTargetInfo
classmethod

Build the info from the target’s text config and the captured layers.

Return the names of the fields whose values differ from other.

class nemo_automodel.components.speculative.dspark.remote.DSparkTargetModelServer(
    target_wrapper: nemo_automodel.components.speculative.dspark.target.HFDSparkTargetModel,
    target_text_config: transformers.PretrainedConfig,
    nccl_port: int,
    host: str = '0.0.0.0'
)

Bases: TargetModelServer

Request-handling logic for a remote DSpark target.

Parameters:

target_wrapper
HFDSparkTargetModel

The frozen target wrapped for DSpark capture.

target_text_config
PretrainedConfig

The config holding the target text stack’s num_hidden_layers, hidden_size and vocab_size.

nccl_port
int

TCP rendezvous port for the NCCL data plane.

host
strDefaults to '0.0.0.0'

Bind/advertise address (rendezvous master for NCCL).

target_info
nemo_automodel.components.speculative.dspark.remote.DSparkTargetModelServer._supervise(
    tensors: dict[str, torch.Tensor]
) -> dict[str, torch.Tensor]

Capture the DSpark features for one decoded request.

Parameters:

tensors
dict[str, torch.Tensor]

Request tensors on the server device: input_ids, attention_mask and loss_mask, each of shape [batch, sequence].

Returns: dict[str, torch.Tensor]

target_hidden_states of shape [batch, sequence, len(target_layer_ids)

nemo_automodel.components.speculative.dspark.remote.DSparkTargetModelServer.handle_model_info(
    _raw: bytes
) -> bytes

Return the served DSparkTargetInfo as JSON.

nemo_automodel.components.speculative.dspark.remote.DSparkTargetModelServer.handle_output_embeddings(
    _raw: bytes
) -> bytes

Return the target lm_head weight (used once to seed the draft head).

nemo_automodel.components.speculative.dspark.remote.DSparkTargetModelServer.post_routes() -> dict[str, tuple[collections.abc.Callable[[bytes], bytes], str]]

Add the output-embedding endpoint the DSpark draft head is seeded from.

class nemo_automodel.components.speculative.dspark.remote.RemoteDSparkTargetModel(
    urls: list[str],
    device: torch.device,
    timeout: int = 120,
    max_retries: int = 3,
    enable_nccl: bool = True
)

DSpark target backend that runs the frozen target on remote servers.

It exposes the surface train_dspark calls on HFDSparkTargetModel (generate_batch, get_input_embeddings, get_output_embeddings). Requests go to the servers round-robin, one at a time per server.

Parameters:

urls
list[str]

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

device
torch.device

Device the returned features are moved to.

timeout
intDefaults to 120

Per-request HTTP timeout in seconds.

max_retries
intDefaults to 3

Retries for idempotent requests on connection errors.

enable_nccl
boolDefaults to True

Whether to use the NCCL data plane when available. A server pairs NCCL with one client only, so pass False when several training ranks share a server.

_pool
nemo_automodel.components.speculative.dspark.remote.RemoteDSparkTargetModel.check_target_info(
    expected: nemo_automodel.components.speculative.dspark.remote.DSparkTargetInfo
) -> None

Raise if any server serves a target that differs from expected.

Parameters:

expected
DSparkTargetInfo

The target properties derived from the training config.

Raises:

  • ValueError: If a server is not a DSpark target server, or reports different properties (naming the server and the mismatched fields).
nemo_automodel.components.speculative.dspark.remote.RemoteDSparkTargetModel.close() -> None

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

nemo_automodel.components.speculative.dspark.remote.RemoteDSparkTargetModel.generate_batch(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor
) -> nemo_automodel.components.speculative.dspark.target.DSparkTargetBatch

Run the remote target once and return the DSpark features.

Parameters:

input_ids
torch.Tensor

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

attention_mask
torch.Tensor

Binary tensor of shape [batch, sequence], with one for valid tokens and zero for padding.

loss_mask
torch.Tensor

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

Returns: DSparkTargetBatch

DSparkTargetBatch with target_hidden_states of shape [batch, sequence,

nemo_automodel.components.speculative.dspark.remote.RemoteDSparkTargetModel.get_input_embeddings() -> types.SimpleNamespace

Fetch the target input-embedding weight, of shape [vocab, hidden] on CPU, as .weight.

nemo_automodel.components.speculative.dspark.remote.RemoteDSparkTargetModel.get_output_embeddings() -> types.SimpleNamespace

Fetch the target lm_head weight, of shape [vocab, hidden] on CPU, as .weight.