mstar.engine.node_dtype

mstar.engine.node_dtype#

Which dtype each node runs in, resolved once at load.

A node’s dtype is its param dtype and its autocast dtype both: the engine manager casts the node to it, and every forward, warm-up and capture runs under autocast in it. None is a node that keeps the dtype it was built in and runs without autocast (NodeSubmodule.disable_autocast).

Functions

node_scope(node_name[, dtype])

Name the node whose forward runs inside, and its resolved dtype, for warn_implicit_downcast.

resolve_model_dtype(model, model_config)

The model-wide dtype: YAML autocast_dtype over get_autocast_dtype.

resolve_node_dtype(submodule, model_dtype)

One node's dtype: None if it opted out, else its own override, else the model's.

warn_implicit_downcast(op, src, dst)

Warn once per node and op that op ran a node's src input in dst rather than the node's dtype.
mstar.engine.node_dtype.node_scope(node_name, dtype=None)[source]#

Name the node whose forward runs inside, and its resolved dtype, for warn_implicit_downcast.

Parameters:

  • node_name (str | None)

  • dtype (dtype | None)

Return type:

Iterator[None]

mstar.engine.node_dtype.resolve_model_dtype(model, model_config)[source]#

The model-wide dtype: YAML autocast_dtype over get_autocast_dtype.

A model returning None pinned its numerics, so YAML may only agree.

Parameters:

  • model (Model)

  • model_config (dict)

Return type:

dtype | None

mstar.engine.node_dtype.resolve_node_dtype(submodule, model_dtype)[source]#

One node’s dtype: None if it opted out, else its own override, else the model’s.

Parameters:
Return type:

dtype | None

mstar.engine.node_dtype.warn_implicit_downcast(op, src, dst)[source]#

Warn once per node and op that op ran a node’s src input in dst rather than the node’s dtype. Silent while dynamo traces, which would graph-break on it; an eager or custom-op caller still reports.

A custom op inside a compiled forward sees no autocast, so a node that autocasts in dst lands here too; that is the dtype it asked for.

Parameters:
Return type:

None