core.ops.loss.backends

Selection of the vocab-parallel cross entropy a model binds.

Module Contents

Functions

select_cross_entropy

Which cross entropy a config asks for, resolved once rather than on every forward.

Data

CrossEntropyTarget

loss = target(logits, labels, tp_group), with logits [s, b, vocab/tp] and the target owning every reduction over tp_group. Targets consume logits: they subtract the row max and exponentiate in place, so a caller that needs them afterwards passes a clone.

API

core.ops.loss.backends.CrossEntropyTarget

None

loss = target(logits, labels, tp_group), with logits [s, b, vocab/tp] and the target owning every reduction over tp_group. Targets consume logits: they subtract the row max and exponentiate in place, so a caller that needs them afterwards passes a clone.

core.ops.loss.backends.select_cross_entropy(

cross_entropy_loss_fusion: bool = False,

cross_entropy_fusion_impl: str = 'native',

cuda_graph_impl: Optional[str] = None,

) → core.ops.loss.backends.CrossEntropyTarget

Which cross entropy a config asks for, resolved once rather than on every forward.

Shared by every provider: the choice depends on the config, not on which backend supplies the rest of the model, which is how it behaved before this was a provider decision.