remotable¶
Giving a model wrapper a remote-execution identity.
A remote run doesn't ship the model — the server already has it loaded. What
travels is a request that names which model to run against: a fully-qualified
model key of the form "import.path.ClassName:model_key". The import path
(resolved with from_import_path) tells the server which
wrapper class to reconstruct; the model-specific suffix (e.g. a HuggingFace repo
id and revision) tells it which checkpoint.
Remotable adds that identity to a model wrapper: to_model_key
mints the key, from_model_key reconstructs a wrapper from one, and the
remote= argument to trace / session routes a run through a
remote (or local-simulation) backend keyed by it. Subclasses supply the two
model-specific halves — _remoteable_model_key and
_remoteable_from_model_key — and may carry per-request state across with
_remoteable_get_env / _remoteable_set_env.
A server also has to know things about a checkpoint it has never loaded, to decide where to put it. Those come in two shapes, and the split is deliberate:
- What the checkpoint is —
describe_checkpointreturns aCheckpointInfo: its size in a given dtype, its parameter count, its config, its revision. One call, because a wrapper that can answer any of these cheaply answers all of them from the same read — a HuggingFace model fetches one config and one Hub record rather than one per question. - What a runtime can do with it —
max_tp_sizestays its own question, and should. How many ways weights can be split is a property of the loader, not of the checkpoint: the same files shard eight ways under transformers tensor parallelism and not at all under something else. Folding it into a description of the checkpoint would make it look like a fact about the files.
Both answer from the key alone, and both have a subclass hook, so a wrapper that can answer cheaply doesn't pay for building the architecture just to count it.
dataclass
¶
CheckpointInfo(size_bytes: Optional[int] = None, n_params: Optional[int] = None, config: Optional[Any] = None, revision: Optional[str] = None)
What a placement decision needs to know about a checkpoint.
Every field is optional and defaults to None: a wrapper answers what it
can, and a caller treats a None as "this wrapper doesn't know" rather
than as a fact. Deliberately not a place for anything runtime-specific —
see the module docstring on why max_tp_size is asked separately.
Bases: Meta
A model wrapper carrying the identity a remote server needs to run it.
See the module docstring for the model-key scheme. Subclasses implement
_remoteable_model_key (the model-specific suffix) and
_remoteable_from_model_key (reconstruct from it), and may override
_remoteable_class when a wrapper should be keyed as another class.
trace(*inputs: Any, backend: Backend | None = None, remote: bool | str = False, blocking: bool = True, job_id: str | None = None, verbose: bool = False, **kwargs: Any) -> Any
session(backend: Backend | None = None, remote: bool | str = False, blocking: bool = True, job_id: str | None = None, verbose: bool = False, tracer_cls: Any = None) -> Any
This model's remote key: "import.path.ClassName:model_key".
The import path names _remoteable_class — the concrete class, or the
canonical class a deprecated alias stands in for, so a model wrapped as
TransformersModel or as its LanguageModel alias produces the one key
the server knows it by. The suffix is _remoteable_model_key. Inverse
of from_model_key.
classmethod
¶
from_model_key(model_key: str, **kwargs: Any) -> Remotable
Reconstruct a model wrapper from a key produced by to_model_key.
Splits off the class import path, resolves it to the wrapper class, and
defers to that class's _remoteable_from_model_key to rebuild from the
model-specific suffix. kwargs are forwarded to the reconstruction.
classmethod
¶
How much memory this checkpoint's weights need, without loading them.
Resolves the wrapper class from the key the same way
from_model_key
does, then asks it. What comes back counts parameters and buffers only —
no activations, no CUDA context, no framework overhead — so a caller
placing a model needs to pad it.
| PARAMETER | DESCRIPTION |
|---|---|
model_key
|
A full key,
TYPE:
|
dtype
|
What the weights will be held in — a torch dtype name, or a
quantization named as a string (see
TYPE:
|
classmethod
¶
What this checkpoint is, without loading it.
One call rather than one per question, because the questions a server asks before placing a model — how big, how many parameters, what architecture, which revision — are all answered from the same metadata, and asking separately meant fetching it repeatedly.
| PARAMETER | DESCRIPTION |
|---|---|
model_key
|
A full key,
TYPE:
|
dtype
|
What the weights will be held in — a torch dtype name, or a
quantization named as a string (see
TYPE:
|
classmethod
¶
The largest number of ranks this checkpoint's weights split into.
None when the model cannot be tensor-parallel at all. Otherwise the
degrees that actually work are the divisors of this number, so a
caller wanting to split a model n ways takes the smallest divisor
>= n — there may be none, and then the model has to be spread some
other way.
How many bytes one weight occupies when held as dtype.
Accepts what torch calls a dtype ("bfloat16", "float32", with or
without a torch. prefix) and what only a quantizer calls one
("int4", "nf4", "fp8"). Fractional for sub-byte formats, so the
result is a float — round the product, not this.
The quantization names come from
QUANTIZATIONS, which is the
same table the loader builds from. Sizing a checkpoint and loading it must
accept exactly the same set: a name only one side knows is a deployment that
is placed and then cannot load, or loads having never been placed. See there
for why these widths are nominal and a caller has to pad.
| RAISES | DESCRIPTION |
|---|---|
ValueError
|
for a name that is neither, rather than defaulting to a plausible width: a wrong guess here silently mis-sizes a deployment. |