Post-Training API 

tuple[PreferenceDataset, PreferenceDataset]

filter(predicate)[source][source]

Return new dataset with filtered pairs.

Parameters:: predicate (Callable[[PreferencePair], bool])
Return type:: PreferenceDataset

shuffle(seed=None)[source][source]

Return new dataset with shuffled pairs.

Parameters:: seed (int | None, default: None)
Return type:: PreferenceDataset

split(train_ratio=0.9, seed=None)[source][source]

Split into train and validation sets.

Parameters:

train_ratio (float, default: 0.9)
seed (int | None, default: None)

Return type:

classmethod from_generations(generator, prompts, reward_fn, num_samples_per_prompt=4, pair_strategy='best_worst', generation_kwargs=None, device='cuda', show_progress=True)[source][source]

Create dataset by generating samples and ranking with reward function.

This is the primary method for creating synthetic preference data.

Parameters:

generator (Module) – Model with a generate method
prompts (Sequence[Any]) – Sequence of conditioning prompts/labels
reward_fn (Callable[[Tensor], Tensor]) – Function that scores samples (higher = better)
num_samples_per_prompt (int, default: 4) – Number of samples to generate per prompt
pair_strategy (Literal['all', 'adjacent', 'best_worst'], default: 'best_worst') – How to create pairs from rankings
generation_kwargs (dict[str, Any] | None, default: None) – Extra kwargs for generator.generate()
device (device | str, default: 'cuda') – Device for generation
show_progress (bool, default: True) – Show progress bar

Return type:

PreferenceDataset

Returns:

PreferenceDataset with generated pairs

classmethod from_json(path)[source][source]

Load from JSON file.

Expected format:

[
    {"chosen": [...], "rejected": [...], "prompt": "...", "margin": 0.8},
    ...
]

Parameters:: path (str | Path)
Return type:: PreferenceDataset

save(path)[source][source]

Save dataset to file.

Parameters:: path (str | Path)
Return type:: None

classmethod load(path)[source][source]

Load dataset from file.

Parameters:: path (str | Path)
Return type:: PreferenceDataset

classmethod __class_getitem__(params)[source]

Parameterizes a generic class.

At least, parameterizing a generic class is the main thing this method does. For example, for some generic class Foo, this is called when we do Foo[int] - there, with cls=Foo and params=int.

However, note that this method is also called when defining generic classes in the first place with class Foo(Generic[T]): ….

class medlatents.post_training.PreferencePair(chosen, rejected, prompt=None, margin=1.0, metadata=<factory>)[source][source]

Bases: object

A single preference comparison between two samples.

Variables:

chosen – The preferred sample (tokens or continuous latents)
rejected – The non-preferred sample
prompt – Optional conditioning information (class label, text, etc.)
margin – Preference strength in [0, 1]. Higher = stronger preference. Default 1.0 for binary preferences, can be soft for uncertain labels.
metadata – Optional dictionary for additional information (source, annotator, etc.)

Example

>>> pair = PreferencePair(
...     chosen=torch.tensor([1, 2, 3]),
...     rejected=torch.tensor([1, 2, 4]),
...     prompt="Generate a brain MRI",
...     margin=0.8,  # Slight preference
... )

Parameters:

chosen (Tensor)
rejected (Tensor)
prompt (Any, default: None)
margin (float, default: 1.0)
metadata (dict[str, Any], default: <factory>)

chosen: torch.Tensor

rejected: torch.Tensor

prompt: Any = None

margin: float = 1.0

metadata: dict[str, Any]

__post_init__()[source][source]

Validate preference pair after initialization.

Return type:: None

to(device)[source][source]

Move tensors to device.

Parameters:: device (device)
Return type:: PreferencePair

__init__(chosen, rejected, prompt=None, margin=1.0, metadata=<factory>)[source]

Parameters:

chosen (Tensor)
rejected (Tensor)
prompt (Any, default: None)
margin (float, default: 1.0)
metadata (dict[str, Any], default: <factory>)

class medlatents.post_training.RankedSamples(samples, prompt=None, scores=None, metadata=<factory>)[source][source]

Bases: object

Multiple samples with a preference ranking.

Useful for generating multiple preference pairs from ranked data. Ranking is from best (index 0) to worst (index -1).

Variables:

samples – List of samples ordered by preference (best first)
prompt – Optional conditioning information
scores – Optional scalar scores for each sample
metadata – Optional dictionary for additional information

Example

>>> ranked = RankedSamples(
...     samples=[best_sample, second_best, worst_sample],
...     scores=[0.9, 0.6, 0.2],
... )
>>> pairs = ranked.to_pairs()  # Creates 3 preference pairs

Parameters:

samples (list[Tensor])
prompt (Any, default: None)
scores (list[float] | None, default: None)
metadata (dict[str, Any], default: <factory>)

samples: list[torch.Tensor]

prompt: Any = None

scores: list[float] | None = None

metadata: dict[str, Any]

__post_init__()[source][source]

Validate ranked samples.

Return type:: None

to_pairs(strategy='all')[source][source]

Convert ranking to preference pairs.

Parameters:

strategy (Literal['all', 'adjacent', 'best_worst'], default: 'all') –

“all”: All pairs where i < j (O(n^2) pairs)
”adjacent”: Only adjacent pairs (O(n) pairs)
”best_worst”: Only best vs worst (1 pair)

Return type:

Returns:

List of PreferencePair objects

__init__(samples, prompt=None, scores=None, metadata=<factory>)[source]

Parameters:

samples (list[Tensor])
prompt (Any, default: None)
scores (list[float] | None, default: None)
metadata (dict[str, Any], default: <factory>)

class medlatents.post_training.StepwisePreference(step_idx, timestep, chosen_state, rejected_state, chosen_score, rejected_score)[source][source]

Bases: object

Preference at a specific denoising step for SPO.

Used by Step-by-step Preference Optimization to capture fine-grained preferences at each noise level.

Variables:

step_idx – Index of the denoising step
timestep – Continuous timestep value in [0, 1]
chosen_state – Preferred intermediate state
rejected_state – Non-preferred intermediate state
chosen_score – Score of chosen state
rejected_score – Score of rejected state

Parameters:

step_idx (int)
timestep (float)
chosen_state (Tensor)
rejected_state (Tensor)
chosen_score (float)
rejected_score (float)

step_idx: int

timestep: float

chosen_state: torch.Tensor

rejected_state: torch.Tensor

chosen_score: float

rejected_score: float

to(device)[source][source]

Move tensors to device.

Parameters:: device (device)
Return type:: StepwisePreference

__init__(step_idx, timestep, chosen_state, rejected_state, chosen_score, rejected_score)[source]

Parameters:

step_idx (int)
timestep (float)
chosen_state (Tensor)
rejected_state (Tensor)
chosen_score (float)
rejected_score (float)

medlatents.post_training.collate_preference_pairs(batch)[source][source]

Collate function for DataLoader with PreferenceDataset.

Returns a dictionary with batched tensors.

Example

>>> loader = DataLoader(dataset, collate_fn=collate_preference_pairs)
>>> for batch in loader:
...     chosen = batch["chosen"]  # [batch_size, seq_len]
...     rejected = batch["rejected"]
...     margins = batch["margin"]

Parameters:: batch (list[PreferencePair])
Return type:: dict[str, Tensor | list[Any]]

class medlatents.post_training.RewardModel(vocab_size=None, in_channels=None, seq_length=256, hidden_size=512, depth=6, num_heads=8, mlp_ratio=4.0, timestep_aware=False, pooling='mean', qk_norm=False, rope_theta=10000.0, dropout=0.1)[source][source]

Trainable reward model that scores samples.

Uses the same transformer architecture as the generative models, but with a scalar output head. This ensures architectural compatibility and enables transfer learning from pretrained generators.

The model can operate in two modes: 1. Discrete: Input is token indices (for MaskGIT, autoreg, D3PM, discrete flow) 2. Continuous: Input is latent vectors (for continuous diffusion/flow)

It optionally accepts timestep conditioning for step-aware scoring (SPO).

Example

>>> # Create reward model for discrete tokens
>>> model = RewardModel(
...     vocab_size=1024,
...     seq_length=256,
...     hidden_size=512,
...     depth=6,
...     num_heads=8,
... )
>>>
>>> # Score samples
>>> scores = model(tokens)  # [batch, 1]
>>>
>>> # Train from preferences
>>> trainer = RewardModelTrainer(model, dataset)
>>> trainer.train(epochs=3)

Parameters:

vocab_size (int | None, default: None)
in_channels (int | None, default: None)
seq_length (int, default: 256)
hidden_size (int, default: 512)
depth (int, default: 6)
num_heads (int, default: 8)
mlp_ratio (float, default: 4.0)
timestep_aware (bool, default: False)
pooling (Literal['mean', 'cls', 'last'], default: 'mean')
qk_norm (bool, default: False)
rope_theta (float, default: 10000.0)
dropout (float, default: 0.1)

__init__(vocab_size=None, in_channels=None, seq_length=256, hidden_size=512, depth=6, num_heads=8, mlp_ratio=4.0, timestep_aware=False, pooling='mean', qk_norm=False, rope_theta=10000.0, dropout=0.1)[source][source]

Initialize reward model.

Parameters:

vocab_size (int | None, default: None) – Vocabulary size for discrete mode (mutually exclusive with in_channels)
in_channels (int | None, default: None) – Input channels for continuous mode
seq_length (int, default: 256) – Maximum sequence length
hidden_size (int, default: 512) – Transformer hidden dimension
depth (int, default: 6) – Number of transformer layers
num_heads (int, default: 8) – Number of attention heads
mlp_ratio (float, default: 4.0) – MLP expansion ratio
timestep_aware (bool, default: False) – If True, accept timestep conditioning
pooling (Literal['mean', 'cls', 'last'], default: 'mean') – How to aggregate sequence for scalar output
qk_norm (bool, default: False) – Use QK LayerNorm for stability
rope_theta (float, default: 10000.0) – RoPE base frequency
dropout (float, default: 0.1) – Dropout rate

forward(x, t=None)[source][source]

Compute reward score for samples.

Parameters:

x (Tensor) – Input samples - Discrete: [batch, seq_len] token indices - Continuous: [batch, seq_len, channels] latent vectors
t (Tensor | None, default: None) – Optional timestep for step-aware scoring [batch]

Return type:

Returns:

Scalar rewards [batch, 1]

compute_preference_loss(chosen, rejected, margin=None, t=None)[source][source]

Compute Bradley-Terry preference loss.

Parameters:

chosen (Tensor) – Preferred samples [batch, …]
rejected (Tensor) – Non-preferred samples [batch, …]
margin (Tensor | None, default: None) – Optional preference strength [batch]
t (Tensor | None, default: None) – Optional timestep for step-aware training

Return type:

dict[str, Tensor]

Returns:

Dict with loss and metrics

T_destination = ~T_destination

add_module(name, module)[source]

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

name (str) – name of the child module. The child module can be accessed from this module using the given name
module (Module) – child module to be added to the module.

Return type:

apply(fn)[source]

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Parameters:: fn (Module -> None) – function to be applied to each submodule
Returns:: self
Return type:: Module

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

bfloat16()[source]

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

buffers(recurse=True)[source]

Return an iterator over module buffers.

Parameters:: recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.
Yields:: torch.Tensor – module buffer
Return type:: Iterator[Tensor]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

call_super_init: bool = False

children()[source]

Return an iterator over immediate children modules.

Yields:: Module – a child module
Return type:: Iterator[Module]

compile(*args, **kwargs)[source]

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()[source]

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

cuda(device=None)[source]

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

double()[source]

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

dump_patches: bool = False

eval()[source]

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:: self
Return type:: Module

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:: str

float()[source]

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

get_buffer(target)[source]

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:: target (str) – The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)
Returns:: The buffer referenced by target
Return type:: torch.Tensor
Raises:: AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:: Any extra state to store in the module’s state_dict
Return type:: object

get_parameter(target)[source]

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:: target (str) – The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)
Returns:: The Parameter referenced by target
Return type:: torch.nn.Parameter
Raises:: AttributeError – If the target string references an invalid path or resolves to something that is not an nn.Parameter

get_submodule(target)[source]

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:: target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)
Returns:: The submodule referenced by target
Return type:: torch.nn.Module
Raises:: AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half()[source]

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

ipu(device=None)[source]

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

load_state_dict(state_dict, strict=True, assign=False)[source]

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

state_dict (dict) – a dict containing parameters and persistent buffers.
strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

modules(remove_duplicate=True)[source]

Return an iterator over all modules in the network.

Parameters:: remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not.
Yields:: Module – a module in the network
Return type:: Iterator[Module]

Note

Duplicate modules are returned only once by default. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)[source]

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

named_buffers(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

prefix (str) – prefix to prepend to all buffer names.
recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.
remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor) – Tuple containing the name and buffer

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

named_children()[source]

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:: (str, Module) – Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

Return type:: Iterator[tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)[source]

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

memo (set[Module] | None, default: None) – a memo to store the set of modules already added to the result
prefix (str, default: '') – a prefix that will be added to the name of the module
remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

prefix (str) – prefix to prepend to all parameter names.
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

parameters(recurse=True)[source]

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:: recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
Yields:: Parameter – module parameter
Return type:: Iterator[Parameter]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

register_backward_hook(hook)[source]

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:: a handle that can be used to remove the added hook by calling handle.remove()
Return type:: torch.utils.hooks.RemovableHandle
Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], tuple[Tensor, ...] | Tensor | None])

register_buffer(name, tensor, persistent=True)[source]

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

name (str) – name of the buffer. The buffer can be accessed from this module using the given name
tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool) – whether the buffer is part of this module’s state_dict.

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)[source]

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Parameters:

hook (Callable) – The user defined hook to be registered.
prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False
with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False
always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)[source]

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Parameters:

hook (Callable) – The user defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False
with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)[source]

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)[source]

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:: a handle that can be used to remove the added hook by calling handle.remove()
Return type:: torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters:: hook (Callable) – Callable hook that will be invoked before loading the state dict.

register_module(name, module)[source]

Alias for add_module().

Parameters:

name (str)
module (Optional[Module])

Return type:

register_parameter(name, param)[source]

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

name (str) – name of the parameter. The parameter can be accessed from this module using the given name
param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Return type:

register_state_dict_post_hook(hook)[source]

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)[source]

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)[source]

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:: requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.
Returns:: self
Return type:: Module

set_extra_state(state)[source]

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:: state (dict) – Extra state from the state_dict
Return type:: None

set_submodule(target, module, strict=False)[source]

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Parameters:

target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)
module (Module) – The module to set the submodule to.
strict (bool, default: False) – If False, the method will replace an existing submodule or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError – If the target string is empty or if module is not an instance of nn.Module.
AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Return type:

share_memory()[source]

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)[source]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)[source]

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)[source]

to(dtype, non_blocking=False)[source]

to(tensor, non_blocking=False)[source]

to(memory_format=torch.channels_last)[source]

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

device (torch.device) – the desired device of the parameters and buffers in this module
dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module
tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)[source]

Move the parameters and buffers to the specified device without copying storage.

Parameters:

device (torch.device) – The desired device of the parameters and buffers in this module.
recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module

train(mode=True)[source]

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:: mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.
Returns:: self
Return type:: Module

type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Parameters:: dst_type (type or string) – the desired type
Returns:: self
Return type:: Module

xpu(device=None)[source]

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

zero_grad(set_to_none=True)[source]

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Parameters:: set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.
Return type:: None

training: bool

class medlatents.post_training.RewardModelTrainer(model, dataset, val_dataset=None, lr=0.0001, batch_size=32, weight_decay=0.01, warmup_ratio=0.1, grad_clip=1.0, device='cuda')[source][source]

Bases: object

Trainer for reward models from preference data.

Handles the training loop with proper logging, checkpointing, and distributed training support.

Example

>>> trainer = RewardModelTrainer(
...     model=reward_model,
...     dataset=preference_dataset,
...     lr=1e-4,
... )
>>> trainer.train(epochs=3)

Parameters:

model (RewardModel)
dataset (PreferenceDataset)
val_dataset (PreferenceDataset | None, default: None)
lr (float, default: 0.0001)
batch_size (int, default: 32)
weight_decay (float, default: 0.01)
warmup_ratio (float, default: 0.1)
grad_clip (float, default: 1.0)
device (device | str, default: 'cuda')

__init__(model, dataset, val_dataset=None, lr=0.0001, batch_size=32, weight_decay=0.01, warmup_ratio=0.1, grad_clip=1.0, device='cuda')[source][source]

Initialize trainer.

Parameters:

model (RewardModel) – RewardModel to train
dataset (PreferenceDataset) – Training preference data
val_dataset (PreferenceDataset | None, default: None) – Optional validation data
lr (float, default: 0.0001) – Learning rate
batch_size (int, default: 32) – Batch size
weight_decay (float, default: 0.01) – Weight decay for AdamW
warmup_ratio (float, default: 0.1) – Fraction of steps for warmup
grad_clip (float, default: 1.0) – Gradient clipping norm
device (device | str, default: 'cuda') – Training device

train(epochs=3, log_every=10, eval_every=100)[source][source]

Train the reward model.

Parameters:

epochs (int, default: 3) – Number of training epochs
log_every (int, default: 10) – Log metrics every N steps
eval_every (int, default: 100) – Run validation every N steps

Return type:

Returns:

Dictionary of training metrics over time

validate()[source][source]

Run validation and return metrics.

Return type:: dict[str, float]

medlatents.post_training.create_reward_model(model_size='small', vocab_size=None, in_channels=None, seq_length=256, timestep_aware=False, **kwargs)[source][source]

Factory function to create reward models with standard configurations.

Parameters:

model_size (Literal['nano', 'small', 'base', 'large'], default: 'small') – Size preset (nano, small, base, large)
vocab_size (int | None, default: None) – For discrete mode
in_channels (int | None, default: None) – For continuous mode
seq_length (int, default: 256) – Maximum sequence length
timestep_aware (bool, default: False) – Enable timestep conditioning for SPO
**kwargs (Any) – Additional arguments to RewardModel

Return type:

RewardModel

Returns:

Configured RewardModel

AI Feedback

class medlatents.post_training.preference.AIPreferenceScorer(score_fn, higher_is_better=True)[source][source]

Bases: object

Generic AI-based preference scoring.

Wraps various scoring backends (reward model, discriminator, external API) into a unified interface.

Example

>>> # From reward model
>>> scorer = AIPreferenceScorer.from_reward_model(rm)
>>>
>>> # Custom scoring function
>>> scorer = AIPreferenceScorer(lambda x: my_score_fn(x))
>>>
>>> # Generate preference pairs
>>> pairs = scorer.generate_pairs(samples, num_pairs=100)

Parameters:

score_fn (Callable[[Tensor], Tensor])
higher_is_better (bool, default: True)

__init__(score_fn, higher_is_better=True)[source][source]

Initialize with a scoring function.

Parameters:

score_fn (Callable[[Tensor], Tensor]) – Function that takes samples and returns scores
higher_is_better (bool, default: True) – If True, higher scores are preferred

classmethod from_reward_model(model)[source][source]

Create scorer from a RewardModel.

Parameters:: model (RewardModel)
Return type:: AIPreferenceScorer

classmethod from_discriminator(discriminator)[source][source]

Create scorer from a discriminator.

Parameters:: discriminator (Module)
Return type:: AIPreferenceScorer

score(samples)[source][source]

Score a batch of samples.

Parameters:: samples (Tensor) – Batch of samples [batch, …]
Return type:: Tensor
Returns:: Scores [batch]

rank(samples)[source][source]

Get ranking of samples (best first).

Parameters:: samples (Sequence[Tensor]) – List of samples
Return type:: list[int]
Returns:: Indices sorted by preference (best first)

create_pairs(samples, prompt=None, strategy='best_worst')[source][source]

Create preference pairs from samples.

Parameters:

samples (Sequence[Tensor]) – Samples to compare
prompt (Any, default: None) – Optional conditioning
strategy (Literal['best_worst', 'adjacent', 'all'], default: 'best_worst') – How to create pairs from ranking

Return type:

Returns:

List of PreferencePair objects

generate_pairs_from_generator(generator, prompts, num_samples_per_prompt=4, strategy='best_worst', generation_kwargs=None, device='cuda')[source][source]

Generate preference pairs by sampling from a generator.

Parameters:

generator (Module) – Model with generate() method
prompts (Sequence[Any]) – Conditioning prompts
num_samples_per_prompt (int, default: 4) – Samples per prompt
strategy (Literal['best_worst', 'adjacent', 'all'], default: 'best_worst') – Pairing strategy
generation_kwargs (dict[str, Any] | None, default: None) – Extra args for generation
device (device | str, default: 'cuda') – Device for generation

Return type:

Returns:

List of PreferencePair objects

class medlatents.post_training.preference.DiscriminatorScorer(discriminator, score_transform='sigmoid')[source][source]

Score samples using a discriminator trained to distinguish real from generated.

This is compatible with your existing TokenCritic and can also wrap any binary classifier.

Example

>>> scorer = DiscriminatorScorer(critic_model)
>>> scores = scorer(generated_samples)  # Higher = more "real"

Parameters:

discriminator (Module)
score_transform (Literal['sigmoid', 'logit', 'raw'], default: 'sigmoid')

__init__(discriminator, score_transform='sigmoid')[source][source]

Initialize discriminator scorer.

Parameters:

discriminator (Module) – Model that outputs realness scores
score_transform (Literal['sigmoid', 'logit', 'raw'], default: 'sigmoid') – How to transform raw outputs - “sigmoid”: Apply sigmoid (if outputs are logits) - “logit”: Keep as logits - “raw”: No transformation

forward(x, **kwargs)[source][source]

Score samples.

Parameters:

x (Tensor) – Input samples [batch, …]
**kwargs (Any) – Additional args for discriminator

Return type:

Returns:

Scores [batch] where higher = more preferred

rank_samples(samples)[source][source]

Rank samples from best to worst.

Parameters:: samples (list[Tensor]) – List of samples to rank
Return type:: list[int]
Returns:: Indices sorted by score (best first)

create_pairs(samples, prompt=None, strategy='best_worst')[source][source]

Create preference pairs from samples using discriminator scores.

Parameters:

samples (list[Tensor]) – List of samples to compare
prompt (Any, default: None) – Optional prompt/conditioning
strategy (Literal['best_worst', 'adjacent', 'all'], default: 'best_worst') – Pairing strategy

Return type:

Returns:

List of PreferencePair objects

T_destination = ~T_destination

add_module(name, module)[source]

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

name (str) – name of the child module. The child module can be accessed from this module using the given name
module (Module) – child module to be added to the module.

Return type:

apply(fn)[source]

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Parameters:: fn (Module -> None) – function to be applied to each submodule
Returns:: self
Return type:: Module

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

bfloat16()[source]

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

buffers(recurse=True)[source]

Return an iterator over module buffers.

Parameters:: recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.
Yields:: torch.Tensor – module buffer
Return type:: Iterator[Tensor]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

call_super_init: bool = False

children()[source]

Return an iterator over immediate children modules.

Yields:: Module – a child module
Return type:: Iterator[Module]

compile(*args, **kwargs)[source]

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()[source]

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

cuda(device=None)[source]

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

double()[source]

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

dump_patches: bool = False

eval()[source]

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:: self
Return type:: Module

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:: str

float()[source]

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

get_buffer(target)[source]

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:: target (str) – The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)
Returns:: The buffer referenced by target
Return type:: torch.Tensor
Raises:: AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:: Any extra state to store in the module’s state_dict
Return type:: object

get_parameter(target)[source]

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:: target (str) – The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)
Returns:: The Parameter referenced by target
Return type:: torch.nn.Parameter
Raises:: AttributeError – If the target string references an invalid path or resolves to something that is not an nn.Parameter

get_submodule(target)[source]

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:: target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)
Returns:: The submodule referenced by target
Return type:: torch.nn.Module
Raises:: AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half()[source]

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

ipu(device=None)[source]

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

load_state_dict(state_dict, strict=True, assign=False)[source]

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

state_dict (dict) – a dict containing parameters and persistent buffers.
strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

modules(remove_duplicate=True)[source]

Return an iterator over all modules in the network.

Parameters:: remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not.
Yields:: Module – a module in the network
Return type:: Iterator[Module]

Note

Duplicate modules are returned only once by default. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)[source]

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

named_buffers(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

prefix (str) – prefix to prepend to all buffer names.
recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.
remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor) – Tuple containing the name and buffer

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

named_children()[source]

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:: (str, Module) – Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

Return type:: Iterator[tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)[source]

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

memo (set[Module] | None, default: None) – a memo to store the set of modules already added to the result
prefix (str, default: '') – a prefix that will be added to the name of the module
remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

prefix (str) – prefix to prepend to all parameter names.
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

parameters(recurse=True)[source]

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:: recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
Yields:: Parameter – module parameter
Return type:: Iterator[Parameter]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

register_backward_hook(hook)[source]

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:: a handle that can be used to remove the added hook by calling handle.remove()
Return type:: torch.utils.hooks.RemovableHandle
Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], tuple[Tensor, ...] | Tensor | None])

register_buffer(name, tensor, persistent=True)[source]

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

name (str) – name of the buffer. The buffer can be accessed from this module using the given name
tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool) – whether the buffer is part of this module’s state_dict.

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)[source]

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Parameters:

hook (Callable) – The user defined hook to be registered.
prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False
with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False
always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)[source]

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Parameters:

hook (Callable) – The user defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False
with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)[source]

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)[source]

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:: a handle that can be used to remove the added hook by calling handle.remove()
Return type:: torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters:: hook (Callable) – Callable hook that will be invoked before loading the state dict.

register_module(name, module)[source]

Alias for add_module().

Parameters:

name (str)
module (Optional[Module])

Return type:

register_parameter(name, param)[source]

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

name (str) – name of the parameter. The parameter can be accessed from this module using the given name
param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Return type:

register_state_dict_post_hook(hook)[source]

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)[source]

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)[source]

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:: requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.
Returns:: self
Return type:: Module

set_extra_state(state)[source]

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:: state (dict) – Extra state from the state_dict
Return type:: None

set_submodule(target, module, strict=False)[source]

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Parameters:

target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)
module (Module) – The module to set the submodule to.
strict (bool, default: False) – If False, the method will replace an existing submodule or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError – If the target string is empty or if module is not an instance of nn.Module.
AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Return type:

share_memory()[source]

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)[source]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)[source]

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)[source]

to(dtype, non_blocking=False)[source]

to(tensor, non_blocking=False)[source]

to(memory_format=torch.channels_last)[source]

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

device (torch.device) – the desired device of the parameters and buffers in this module
dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module
tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)[source]

Move the parameters and buffers to the specified device without copying storage.

Parameters:

device (torch.device) – The desired device of the parameters and buffers in this module.
recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module

train(mode=True)[source]

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:: mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.
Returns:: self
Return type:: Module

type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Parameters:: dst_type (type or string) – the desired type
Returns:: self
Return type:: Module

xpu(device=None)[source]

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

zero_grad(set_to_none=True)[source]

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Parameters:: set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.
Return type:: None

training: bool

class medlatents.post_training.preference.StepAwarePreferenceModel(base_model=None, vocab_size=None, in_channels=None, seq_length=256, hidden_size=384, depth=6, num_heads=6)[source][source]

Timestep-aware preference model for Step-by-step Preference Optimization (SPO).

This model scores intermediate denoising states, allowing SPO to learn fine-grained preferences at each noise level.

The key insight is that preferences differ at different timesteps: - Early steps (high noise): Layout/composition preferences - Late steps (low noise): Fine detail preferences

Example

>>> step_model = StepAwarePreferenceModel(
...     base_model=reward_model,  # Must be timestep_aware=True
... )
>>> scores = step_model(noisy_samples, t=timesteps)

Parameters:

base_model (RewardModel | None, default: None)
vocab_size (int | None, default: None)
in_channels (int | None, default: None)
seq_length (int, default: 256)
hidden_size (int, default: 384)
depth (int, default: 6)
num_heads (int, default: 6)

__init__(base_model=None, vocab_size=None, in_channels=None, seq_length=256, hidden_size=384, depth=6, num_heads=6)[source][source]

Initialize step-aware preference model.

Either provide a base_model (must have timestep_aware=True) or specify architecture parameters to create a new model.

Parameters:

base_model (RewardModel | None, default: None) – Pre-existing timestep-aware RewardModel
vocab_size (int | None, default: None) – For creating new model (discrete)
in_channels (int | None, default: None) – For creating new model (continuous)
seq_length (int, default: 256) – Sequence length
hidden_size (int, default: 384) – Hidden dimension
depth (int, default: 6) – Number of layers
num_heads (int, default: 6) – Attention heads

forward(x, t)[source][source]

Score samples at given timesteps.

Parameters:

x (Tensor) – Noisy samples [batch, …]
t (Tensor) – Timesteps [batch] in [0, 1]

Return type:

Returns:

Scores [batch]

select_winner_loser(candidates, t)[source][source]

Select best and worst candidates for SPO training.

Parameters:

candidates (Tensor) – Candidate samples [num_candidates, batch, …]
t (Tensor) – Timestep [batch] or scalar

Return type:

tuple[Tensor, Tensor, Tensor, Tensor]

Returns:

(chosen, rejected, chosen_scores, rejected_scores)

sample_random_candidate(candidates)[source][source]

Randomly select one candidate per batch item (for next SPO step).

Parameters:: candidates (Tensor) – [num_candidates, batch, …]
Return type:: Tensor
Returns:: Selected samples [batch, …]

T_destination = ~T_destination

add_module(name, module)[source]

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

name (str) – name of the child module. The child module can be accessed from this module using the given name
module (Module) – child module to be added to the module.

Return type:

apply(fn)[source]

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Parameters:: fn (Module -> None) – function to be applied to each submodule
Returns:: self
Return type:: Module

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

bfloat16()[source]

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

buffers(recurse=True)[source]

Return an iterator over module buffers.

Parameters:: recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.
Yields:: torch.Tensor – module buffer
Return type:: Iterator[Tensor]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

call_super_init: bool = False

children()[source]

Return an iterator over immediate children modules.

Yields:: Module – a child module
Return type:: Iterator[Module]

compile(*args, **kwargs)[source]

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()[source]

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

cuda(device=None)[source]

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

double()[source]

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

dump_patches: bool = False

eval()[source]

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:: self
Return type:: Module

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:: str

float()[source]

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

get_buffer(target)[source]

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:: target (str) – The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)
Returns:: The buffer referenced by target
Return type:: torch.Tensor
Raises:: AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:: Any extra state to store in the module’s state_dict
Return type:: object

get_parameter(target)[source]

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:: target (str) – The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)
Returns:: The Parameter referenced by target
Return type:: torch.nn.Parameter
Raises:: AttributeError – If the target string references an invalid path or resolves to something that is not an nn.Parameter

get_submodule(target)[source]

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:: target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)
Returns:: The submodule referenced by target
Return type:: torch.nn.Module
Raises:: AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half()[source]

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

ipu(device=None)[source]

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

load_state_dict(state_dict, strict=True, assign=False)[source]

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

state_dict (dict) – a dict containing parameters and persistent buffers.
strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

modules(remove_duplicate=True)[source]

Return an iterator over all modules in the network.

Parameters:: remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not.
Yields:: Module – a module in the network
Return type:: Iterator[Module]

Note

Duplicate modules are returned only once by default. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)[source]

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

named_buffers(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

prefix (str) – prefix to prepend to all buffer names.
recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.
remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor) – Tuple containing the name and buffer

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

named_children()[source]

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:: (str, Module) – Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

Return type:: Iterator[tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)[source]

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

memo (set[Module] | None, default: None) – a memo to store the set of modules already added to the result
prefix (str, default: '') – a prefix that will be added to the name of the module
remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

prefix (str) – prefix to prepend to all parameter names.
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

parameters(recurse=True)[source]

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:: recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
Yields:: Parameter – module parameter
Return type:: Iterator[Parameter]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

register_backward_hook(hook)[source]

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:: a handle that can be used to remove the added hook by calling handle.remove()
Return type:: torch.utils.hooks.RemovableHandle
Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], tuple[Tensor, ...] | Tensor | None])

register_buffer(name, tensor, persistent=True)[source]

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

name (str) – name of the buffer. The buffer can be accessed from this module using the given name
tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool) – whether the buffer is part of this module’s state_dict.

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)[source]

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Parameters:

hook (Callable) – The user defined hook to be registered.
prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False
with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False
always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)[source]

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Parameters:

hook (Callable) – The user defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False
with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)[source]

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)[source]

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:: a handle that can be used to remove the added hook by calling handle.remove()
Return type:: torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters:: hook (Callable) – Callable hook that will be invoked before loading the state dict.

register_module(name, module)[source]

Alias for add_module().

Parameters:

name (str)
module (Optional[Module])

Return type:

register_parameter(name, param)[source]

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

name (str) – name of the parameter. The parameter can be accessed from this module using the given name
param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Return type:

register_state_dict_post_hook(hook)[source]

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)[source]

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)[source]

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:: requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.
Returns:: self
Return type:: Module

set_extra_state(state)[source]

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:: state (dict) – Extra state from the state_dict
Return type:: None

set_submodule(target, module, strict=False)[source]

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Parameters:

target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)
module (Module) – The module to set the submodule to.
strict (bool, default: False) – If False, the method will replace an existing submodule or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError – If the target string is empty or if module is not an instance of nn.Module.
AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Return type:

share_memory()[source]

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)[source]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)[source]

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)[source]

to(dtype, non_blocking=False)[source]

to(tensor, non_blocking=False)[source]

to(memory_format=torch.channels_last)[source]

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

device (torch.device) – the desired device of the parameters and buffers in this module
dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module
tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)[source]

Move the parameters and buffers to the specified device without copying storage.

Parameters:

device (torch.device) – The desired device of the parameters and buffers in this module.
recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module

train(mode=True)[source]

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:: mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.
Returns:: self
Return type:: Module

type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Parameters:: dst_type (type or string) – the desired type
Returns:: self
Return type:: Module

xpu(device=None)[source]

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

zero_grad(set_to_none=True)[source]

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Parameters:: set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.
Return type:: None

training: bool

Direct Preference Optimization (DPO)

class medlatents.post_training.BaseDPOTrainer(model, ref_model, config, accelerator=None)[source][source]

Bases: abc.ABC

Abstract base class for DPO trainers.

Provides common infrastructure for all model types: - Accelerator setup with FSDP support - EMA for stable training - Logging and checkpointing - Reference model management

Subclasses must implement: - compute_logprobs: Model-specific log-probability computation - prepare_batch: Model-specific batch preprocessing

Example

>>> class MyDPOTrainer(BaseDPOTrainer):
...     def compute_logprobs(self, model, batch):
...         # Model-specific implementation
...         return log_probs
...
>>> trainer = MyDPOTrainer(model, ref_model, config)
>>> trainer.train(dataset)

Parameters:

model (Module)
ref_model (Module | None)
config (PostTrainingConfig)
accelerator (Accelerator | None, default: None)

__init__(model, ref_model, config, accelerator=None)[source][source]

Initialize DPO trainer.

Parameters:

model (Module) – Model to train
ref_model (Module | None) – Reference model (frozen). If None, uses reference_free mode.
config (PostTrainingConfig) – Training configuration
accelerator (Accelerator | None, default: None) – Optional pre-configured accelerator

abstractmethod compute_logprobs(model, samples, **kwargs)[source][source]

Compute log-probabilities for samples.

This is model-type specific: - Autoregressive: Sum of log P(x_i | x_{<i}) - MaskGIT: Pseudo-likelihood via masked prediction - Diffusion: Negative ELBO (Monte Carlo estimate) - Flow: Trajectory-based likelihood

Parameters:

model (Module) – Model to evaluate
samples (Tensor) – Token sequences [batch, seq_len]
**kwargs (Any) – Model-specific arguments

Return type:

Returns:

Log-probabilities [batch]

prepare_batch(batch)[source][source]

Prepare batch for DPO training.

Parameters:: batch (dict[str, Any]) – Collated batch from DataLoader
Return type:: tuple[Tensor, Tensor, Tensor, dict[str, Any]]
Returns:: Tuple of (chosen, rejected, margin, extra_kwargs)

train_step(batch)[source][source]

Execute a single training step.

Parameters:: batch (dict[str, Any]) – Collated preference batch
Return type:: dict[str, float]
Returns:: Dictionary of metrics

train(dataset, val_dataset=None)[source][source]

Run full training loop.

Parameters:

dataset (PreferenceDataset) – Training preference dataset
val_dataset (PreferenceDataset | None, default: None) – Optional validation dataset

Return type:

Returns:

Dictionary of training history

validate(val_loader)[source][source]

Run validation.

Parameters:: val_loader (DataLoader) – Validation dataloader
Return type:: dict[str, float]
Returns:: Validation metrics

save_checkpoint(name='latest')[source][source]

Save training checkpoint.

Parameters:: name (str, default: 'latest') – Checkpoint name (e.g., “latest”, “best”, “step_1000”)
Return type:: None

load_checkpoint(path)[source][source]

Load training checkpoint.

Parameters:: path (str) – Path to checkpoint file
Return type:: None

class medlatents.post_training.DPOLoss(beta=0.1, loss_type='sigmoid', label_smoothing=0.0, reference_free=False, kto_desirable_weight=1.0, kto_undesirable_weight=1.0)[source][source]

Flexible DPO loss with multiple variants.

Supports: - sigmoid: Standard DPO with sigmoid loss (Rafailov et al.) - hinge: Margin-based loss for more aggressive optimization - ipo: Identity Preference Optimization (no sigmoid) - kto: Kahneman-Tversky Optimization (asymmetric) - disco: Distributionally robust DPO

The DPO objective is:: L = -log(sigmoid(β * (r_w - r_l)))
where r_w and r_l are implicit rewards defined as:: r(x) = β * log(π_θ(x) / π_ref(x))

Example

>>> loss_fn = DPOLoss(beta=0.1, loss_type="sigmoid")
>>> output = loss_fn(
...     chosen_logprobs=chosen_lp,
...     rejected_logprobs=rejected_lp,
...     ref_chosen_logprobs=ref_chosen_lp,
...     ref_rejected_logprobs=ref_rejected_lp,
... )
>>> loss = output.loss

Parameters:

beta (float, default: 0.1)
loss_type (Literal['sigmoid', 'hinge', 'ipo', 'kto', 'disco'], default: 'sigmoid')
label_smoothing (float, default: 0.0)
reference_free (bool, default: False)
kto_desirable_weight (float, default: 1.0)
kto_undesirable_weight (float, default: 1.0)

__init__(beta=0.1, loss_type='sigmoid', label_smoothing=0.0, reference_free=False, kto_desirable_weight=1.0, kto_undesirable_weight=1.0)[source][source]

Initialize DPO loss.

Parameters:

beta (float, default: 0.1) – KL penalty coefficient (controls deviation from reference)
loss_type (Literal['sigmoid', 'hinge', 'ipo', 'kto', 'disco'], default: 'sigmoid') – Type of preference loss to use
label_smoothing (float, default: 0.0) – Soft labels for robustness
reference_free (bool, default: False) – If True, don’t use reference model
kto_desirable_weight (float, default: 1.0) – Weight for KTO desirable term
kto_undesirable_weight (float, default: 1.0) – Weight for KTO undesirable term

forward(chosen_logprobs, rejected_logprobs, ref_chosen_logprobs=None, ref_rejected_logprobs=None, margin=None)[source][source]

Compute DPO loss.

Parameters:

chosen_logprobs (Tensor) – Log probs of policy on chosen [batch]
rejected_logprobs (Tensor) – Log probs of policy on rejected [batch]
ref_chosen_logprobs (Tensor | None, default: None) – Log probs of reference on chosen [batch]
ref_rejected_logprobs (Tensor | None, default: None) – Log probs of reference on rejected [batch]
margin (Tensor | None, default: None) – Optional preference margin [batch]

Return type:

DPOOutput

Returns:

DPOOutput with loss and metrics

T_destination = ~T_destination

add_module(name, module)[source]

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

name (str) – name of the child module. The child module can be accessed from this module using the given name
module (Module) – child module to be added to the module.

Return type:

apply(fn)[source]

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Parameters:: fn (Module -> None) – function to be applied to each submodule
Returns:: self
Return type:: Module

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

bfloat16()[source]

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

buffers(recurse=True)[source]

Return an iterator over module buffers.

Parameters:: recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.
Yields:: torch.Tensor – module buffer
Return type:: Iterator[Tensor]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

call_super_init: bool = False

children()[source]

Return an iterator over immediate children modules.

Yields:: Module – a child module
Return type:: Iterator[Module]

compile(*args, **kwargs)[source]

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()[source]

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

cuda(device=None)[source]

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

double()[source]

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

dump_patches: bool = False

eval()[source]

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:: self
Return type:: Module

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:: str

float()[source]

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

get_buffer(target)[source]

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:: target (str) – The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)
Returns:: The buffer referenced by target
Return type:: torch.Tensor
Raises:: AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:: Any extra state to store in the module’s state_dict
Return type:: object

get_parameter(target)[source]

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:: target (str) – The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)
Returns:: The Parameter referenced by target
Return type:: torch.nn.Parameter
Raises:: AttributeError – If the target string references an invalid path or resolves to something that is not an nn.Parameter

get_submodule(target)[source]

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:: target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)
Returns:: The submodule referenced by target
Return type:: torch.nn.Module
Raises:: AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half()[source]

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

ipu(device=None)[source]

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

load_state_dict(state_dict, strict=True, assign=False)[source]

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

state_dict (dict) – a dict containing parameters and persistent buffers.
strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

modules(remove_duplicate=True)[source]

Return an iterator over all modules in the network.

Parameters:: remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not.
Yields:: Module – a module in the network
Return type:: Iterator[Module]

Note

Duplicate modules are returned only once by default. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)[source]

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

named_buffers(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

prefix (str) – prefix to prepend to all buffer names.
recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.
remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor) – Tuple containing the name and buffer

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

named_children()[source]

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:: (str, Module) – Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

Return type:: Iterator[tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)[source]

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

memo (set[Module] | None, default: None) – a memo to store the set of modules already added to the result
prefix (str, default: '') – a prefix that will be added to the name of the module
remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

prefix (str) – prefix to prepend to all parameter names.
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

parameters(recurse=True)[source]

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:: recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
Yields:: Parameter – module parameter
Return type:: Iterator[Parameter]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

register_backward_hook(hook)[source]

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:: a handle that can be used to remove the added hook by calling handle.remove()
Return type:: torch.utils.hooks.RemovableHandle
Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], tuple[Tensor, ...] | Tensor | None])

register_buffer(name, tensor, persistent=True)[source]

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

name (str) – name of the buffer. The buffer can be accessed from this module using the given name
tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool) – whether the buffer is part of this module’s state_dict.

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)[source]

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Parameters:

hook (Callable) – The user defined hook to be registered.
prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False
with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False
always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)[source]

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Parameters:

hook (Callable) – The user defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False
with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)[source]

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)[source]

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:: a handle that can be used to remove the added hook by calling handle.remove()
Return type:: torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters:: hook (Callable) – Callable hook that will be invoked before loading the state dict.

register_module(name, module)[source]

Alias for add_module().

Parameters:

name (str)
module (Optional[Module])

Return type:

register_parameter(name, param)[source]

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

name (str) – name of the parameter. The parameter can be accessed from this module using the given name
param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Return type:

register_state_dict_post_hook(hook)[source]

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)[source]

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)[source]

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:: requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.
Returns:: self
Return type:: Module

set_extra_state(state)[source]

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:: state (dict) – Extra state from the state_dict
Return type:: None

set_submodule(target, module, strict=False)[source]

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Parameters:

target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)
module (Module) – The module to set the submodule to.
strict (bool, default: False) – If False, the method will replace an existing submodule or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError – If the target string is empty or if module is not an instance of nn.Module.
AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Return type:

share_memory()[source]

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)[source]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)[source]

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)[source]

to(dtype, non_blocking=False)[source]

to(tensor, non_blocking=False)[source]

to(memory_format=torch.channels_last)[source]

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

device (torch.device) – the desired device of the parameters and buffers in this module
dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module
tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)[source]

Move the parameters and buffers to the specified device without copying storage.

Parameters:

device (torch.device) – The desired device of the parameters and buffers in this module.
recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module

train(mode=True)[source]

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:: mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.
Returns:: self
Return type:: Module

type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Parameters:: dst_type (type or string) – the desired type
Returns:: self
Return type:: Module

xpu(device=None)[source]

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

zero_grad(set_to_none=True)[source]

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Parameters:: set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.
Return type:: None

training: bool

class medlatents.post_training.DPOOutput(loss, chosen_rewards, rejected_rewards, accuracy, reward_margin)[source][source]

Bases: object

Output from DPO loss computation.

Variables:

loss – The total DPO loss
chosen_rewards – Implicit rewards for chosen samples (log-prob ratio)
rejected_rewards – Implicit rewards for rejected samples
accuracy – Fraction where chosen_reward > rejected_reward
reward_margin – Mean (chosen_reward - rejected_reward)

Parameters:

loss (Tensor)
chosen_rewards (Tensor)
rejected_rewards (Tensor)
accuracy (Tensor)
reward_margin (Tensor)

loss: torch.Tensor

chosen_rewards: torch.Tensor

rejected_rewards: torch.Tensor

accuracy: torch.Tensor

reward_margin: torch.Tensor

to_dict()[source][source]

Convert to logging dictionary.

Return type:: dict[str, float]

__init__(loss, chosen_rewards, rejected_rewards, accuracy, reward_margin)[source]

Parameters:

loss (Tensor)
chosen_rewards (Tensor)
rejected_rewards (Tensor)
accuracy (Tensor)
reward_margin (Tensor)

Discrete models (autoregressive, MaskGIT):

class medlatents.post_training.AutoregressiveDPOTrainer(model, ref_model, config, accelerator=None, vocab_size=None)[source][source]

Bases: medlatents.post_training.dpo.base.BaseDPOTrainer

DPO trainer for autoregressive discrete models.

For autoregressive models, the log-probability is computed as:: log p(x) = sum_{i=1}^L log p(x_i | x_{<i})

This is the standard approach used in language model DPO.

Example

>>> from medlatents.autoregressive import AutoregressiveTransformer
>>> model = AutoregressiveTransformer(vocab_size=1024, ...)
>>> ref_model = copy.deepcopy(model)
>>> trainer = AutoregressiveDPOTrainer(model, ref_model, config)
>>> trainer.train(preference_dataset)

Parameters:

model (Module)
ref_model (Module | None)
config (PostTrainingConfig)
accelerator (Accelerator | None, default: None)
vocab_size (int | None, default: None)

__init__(model, ref_model, config, accelerator=None, vocab_size=None)[source][source]

Initialize autoregressive DPO trainer.

Parameters:

model (Module) – Autoregressive model to train
ref_model (Module | None) – Frozen reference model
config (PostTrainingConfig) – Training configuration
accelerator (Accelerator | None, default: None) – Optional pre-configured accelerator
vocab_size (int | None, default: None) – Vocabulary size (inferred from model if not provided)

compute_logprobs(model, samples, **kwargs)[source][source]

Compute autoregressive log-probabilities.

For a sequence x = [x_1, x_2, …, x_L], computes:: log p(x) = sum_{i=1}^{L-1} log p(x_{i+1} | x_1, …, x_i)

Parameters:

model (Module) – Autoregressive model
samples (Tensor) – Token sequences [batch, seq_len]
**kwargs (Any) – Additional model arguments

Return type:

Returns:

Log-probabilities [batch]

load_checkpoint(path)[source]

Load training checkpoint.

Parameters:: path (str) – Path to checkpoint file
Return type:: None

prepare_batch(batch)[source]

Prepare batch for DPO training.

Parameters:: batch (dict[str, Any]) – Collated batch from DataLoader
Return type:: tuple[Tensor, Tensor, Tensor, dict[str, Any]]
Returns:: Tuple of (chosen, rejected, margin, extra_kwargs)

save_checkpoint(name='latest')[source]

Save training checkpoint.

Parameters:: name (str, default: 'latest') – Checkpoint name (e.g., “latest”, “best”, “step_1000”)
Return type:: None

train(dataset, val_dataset=None)[source]

Run full training loop.

Parameters:

dataset (PreferenceDataset) – Training preference dataset
val_dataset (PreferenceDataset | None, default: None) – Optional validation dataset

Return type:

Returns:

Dictionary of training history

train_step(batch)[source]

Execute a single training step.

Parameters:: batch (dict[str, Any]) – Collated preference batch
Return type:: dict[str, float]
Returns:: Dictionary of metrics

validate(val_loader)[source]

Run validation.

Parameters:: val_loader (DataLoader) – Validation dataloader
Return type:: dict[str, float]
Returns:: Validation metrics

class medlatents.post_training.MaskGITDPOTrainer(model, ref_model, config, accelerator=None, vocab_size=None, likelihood_strategy='parallel', num_mask_samples=16)[source][source]

Bases: medlatents.post_training.dpo.base.BaseDPOTrainer

DPO trainer for MaskGIT bidirectional models.

For bidirectional models, we use pseudo-likelihood:: log p(x) ≈ sum_{i=1}^L log p(x_i | x_{i})

where x_{i} denotes all tokens except position i.

This is computed efficiently by masking each position and predicting it from the context. We can use different strategies:

“full”: Mask each position independently (L forward passes, expensive)
“random”: Random subset of positions (cheaper approximation)
“parallel”: Single forward with random masking (training-style)

Example

>>> from medlatents.maskgit import MaskGIT
>>> model = MaskGIT(vocab_size=1024, ...)
>>> ref_model = copy.deepcopy(model)
>>> trainer = MaskGITDPOTrainer(model, ref_model, config)
>>> trainer.train(preference_dataset)

Parameters:

model (Module)
ref_model (Module | None)
config (PostTrainingConfig)
accelerator (Accelerator | None, default: None)
vocab_size (int | None, default: None)
likelihood_strategy (Literal['full', 'random', 'parallel'], default: 'parallel')
num_mask_samples (int, default: 16)

__init__(model, ref_model, config, accelerator=None, vocab_size=None, likelihood_strategy='parallel', num_mask_samples=16)[source][source]

Initialize MaskGIT DPO trainer.

Parameters:

model (Module) – MaskGIT model to train
ref_model (Module | None) – Frozen reference model
config (PostTrainingConfig) – Training configuration
accelerator (Accelerator | None, default: None) – Optional pre-configured accelerator
vocab_size (int | None, default: None) – Vocabulary size (inferred from model if not provided)
likelihood_strategy (Literal['full', 'random', 'parallel'], default: 'parallel') – How to compute pseudo-likelihood
num_mask_samples (int, default: 16) – Number of positions to sample for “random” strategy

compute_logprobs(model, samples, **kwargs)[source][source]

Compute pseudo-likelihood for MaskGIT.

Uses the specified strategy to estimate log p(x).

Parameters:

model (Module) – MaskGIT model
samples (Tensor) – Token sequences [batch, seq_len]
**kwargs (Any) – Additional model arguments

Return type:

Returns:

Log-probabilities [batch]

load_checkpoint(path)[source]

Load training checkpoint.

Parameters:: path (str) – Path to checkpoint file
Return type:: None

prepare_batch(batch)[source]

Prepare batch for DPO training.

Parameters:: batch (dict[str, Any]) – Collated batch from DataLoader
Return type:: tuple[Tensor, Tensor, Tensor, dict[str, Any]]
Returns:: Tuple of (chosen, rejected, margin, extra_kwargs)

save_checkpoint(name='latest')[source]

Save training checkpoint.

Parameters:: name (str, default: 'latest') – Checkpoint name (e.g., “latest”, “best”, “step_1000”)
Return type:: None

train(dataset, val_dataset=None)[source]

Run full training loop.

Parameters:

dataset (PreferenceDataset) – Training preference dataset
val_dataset (PreferenceDataset | None, default: None) – Optional validation dataset

Return type:

Returns:

Dictionary of training history

train_step(batch)[source]

Execute a single training step.

Parameters:: batch (dict[str, Any]) – Collated preference batch
Return type:: dict[str, float]
Returns:: Dictionary of metrics

validate(val_loader)[source]

Run validation.

Parameters:: val_loader (DataLoader) – Validation dataloader
Return type:: dict[str, float]
Returns:: Validation metrics

class medlatents.post_training.DiscreteDPOTrainer(model, ref_model, config, accelerator=None, model_type=None, vocab_size=None, **kwargs)[source][source]

Bases: medlatents.post_training.dpo.base.BaseDPOTrainer

Unified DPO trainer that auto-detects discrete model type.

This is a convenience class that automatically selects the appropriate training strategy based on the model type.

Example

>>> # Works with any discrete model
>>> trainer = DiscreteDPOTrainer(model, ref_model, config)
>>> trainer.train(dataset)

Parameters:

model (Module)
ref_model (Module | None)
config (PostTrainingConfig)
accelerator (Accelerator | None, default: None)
model_type (Optional[Literal['autoreg', 'maskgit']], default: None)
vocab_size (int | None, default: None)
kwargs (Any)

__init__(model, ref_model, config, accelerator=None, model_type=None, vocab_size=None, **kwargs)[source][source]

Initialize discrete DPO trainer.

Parameters:

model (Module) – Discrete model to train
ref_model (Module | None) – Frozen reference model
config (PostTrainingConfig) – Training configuration
accelerator (Accelerator | None, default: None) – Optional pre-configured accelerator
model_type (Optional[Literal['autoreg', 'maskgit']], default: None) – Model type (“autoreg” or “maskgit”). Auto-detected if None.
vocab_size (int | None, default: None) – Vocabulary size
**kwargs (Any) – Additional arguments passed to specific trainer

compute_logprobs(model, samples, **kwargs)[source][source]

Delegate to internal trainer.

Parameters:

model (Module)
samples (Tensor)
kwargs (Any)

Return type:

train(*args, **kwargs)[source][source]: Delegate training to internal trainer.

validate(*args, **kwargs)[source][source]: Delegate validation to internal trainer.

prepare_batch(batch)[source]

Prepare batch for DPO training.

Parameters:: batch (dict[str, Any]) – Collated batch from DataLoader
Return type:: tuple[Tensor, Tensor, Tensor, dict[str, Any]]
Returns:: Tuple of (chosen, rejected, margin, extra_kwargs)

save_checkpoint(*args, **kwargs)[source][source]: Delegate checkpoint saving to internal trainer.

train_step(batch)[source]

Execute a single training step.

Parameters:: batch (dict[str, Any]) – Collated preference batch
Return type:: dict[str, float]
Returns:: Dictionary of metrics

load_checkpoint(*args, **kwargs)[source][source]: Delegate checkpoint loading to internal trainer.

Diffusion models (D3PM):

class medlatents.post_training.D3PMDPOTrainer(model, ref_model, config, d3pm, accelerator=None, num_timestep_samples=10, elbo_mode='monte_carlo')[source][source]

Bases: medlatents.post_training.dpo.base.BaseDPOTrainer

DPO trainer for D3PM discrete diffusion models.

For D3PM models, we compute log-probability using the ELBO:: log p(x_0) ≥ -ELBO = -E_t[KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))]

The ELBO provides a lower bound on the log-likelihood, which we can use for DPO optimization.

We support two modes: 1. Monte Carlo ELBO: Sample timesteps and average 2. Full ELBO: Compute over all timesteps (expensive but exact)

Example

>>> from medlatents.diffusion import D3PM
>>> d3pm = D3PM(num_classes=1024, num_timesteps=1000)
>>> model = DiscreteDiT(vocab_size=1024, ...)
>>> ref_model = copy.deepcopy(model)
>>> trainer = D3PMDPOTrainer(model, ref_model, config, d3pm=d3pm)
>>> trainer.train(preference_dataset)

Parameters:

model (Module)
ref_model (Module | None)
config (PostTrainingConfig)
d3pm (D3PM)
accelerator (Accelerator | None, default: None)
num_timestep_samples (int, default: 10)
elbo_mode (Literal['monte_carlo', 'full'], default: 'monte_carlo')

__init__(model, ref_model, config, d3pm, accelerator=None, num_timestep_samples=10, elbo_mode='monte_carlo')[source][source]

Initialize D3PM DPO trainer.

Parameters:

model (Module) – D3PM backbone model to train
ref_model (Module | None) – Frozen reference model
config (PostTrainingConfig) – Training configuration
d3pm (D3PM) – D3PM diffusion process
accelerator (Accelerator | None, default: None) – Optional pre-configured accelerator
num_timestep_samples (int, default: 10) – Number of timesteps to sample for Monte Carlo ELBO
elbo_mode (Literal['monte_carlo', 'full'], default: 'monte_carlo') – How to compute ELBO (“monte_carlo” or “full”)

compute_logprobs(model, samples, **kwargs)[source][source]

Compute ELBO-based log-probability for D3PM.

The ELBO is computed as:: ELBO = E_t[KL(q(x_{t-1}|x_t,x_0) || p_θ(x_{t-1}|x_t))]

We return -ELBO as a surrogate for log p(x).

Parameters:

model (Module) – D3PM backbone model
samples (Tensor) – Clean token sequences [batch, seq_len]
**kwargs (Any) – Additional model arguments

Return type:

Returns:

Log-probabilities (negative ELBO) [batch]

load_checkpoint(path)[source]

Load training checkpoint.

Parameters:: path (str) – Path to checkpoint file
Return type:: None

prepare_batch(batch)[source]

Prepare batch for DPO training.

Parameters:: batch (dict[str, Any]) – Collated batch from DataLoader
Return type:: tuple[Tensor, Tensor, Tensor, dict[str, Any]]
Returns:: Tuple of (chosen, rejected, margin, extra_kwargs)

save_checkpoint(name='latest')[source]

Save training checkpoint.

Parameters:: name (str, default: 'latest') – Checkpoint name (e.g., “latest”, “best”, “step_1000”)
Return type:: None

train(dataset, val_dataset=None)[source]

Run full training loop.

Parameters:

dataset (PreferenceDataset) – Training preference dataset
val_dataset (PreferenceDataset | None, default: None) – Optional validation dataset

Return type:

Returns:

Dictionary of training history

train_step(batch)[source]

Execute a single training step.

Parameters:: batch (dict[str, Any]) – Collated preference batch
Return type:: dict[str, float]
Returns:: Dictionary of metrics

validate(val_loader)[source]

Run validation.

Parameters:: val_loader (DataLoader) – Validation dataloader
Return type:: dict[str, float]
Returns:: Validation metrics

class medlatents.post_training.DiffusionDPOTrainer(model, ref_model, config, diffusion, accelerator=None, num_timestep_samples=10, elbo_mode='monte_carlo')[source][source]

Bases: medlatents.post_training.dpo.base.BaseDPOTrainer

Unified DPO trainer for diffusion models.

This is a convenience class that handles both D3PM discrete diffusion and could be extended for continuous diffusion.

For now, it primarily wraps D3PM but provides a cleaner interface.

Example

>>> trainer = DiffusionDPOTrainer(
...     model=backbone_model,
...     ref_model=ref_backbone,
...     config=config,
...     diffusion=d3pm,
... )
>>> trainer.train(preference_dataset)

Parameters:

model (Module)
ref_model (Module | None)
config (PostTrainingConfig)
diffusion (D3PM)
accelerator (Accelerator | None, default: None)
num_timestep_samples (int, default: 10)
elbo_mode (Literal['monte_carlo', 'full'], default: 'monte_carlo')

__init__(model, ref_model, config, diffusion, accelerator=None, num_timestep_samples=10, elbo_mode='monte_carlo')[source][source]

Initialize diffusion DPO trainer.

Parameters:

model (Module) – Diffusion backbone model to train
ref_model (Module | None) – Frozen reference model
config (PostTrainingConfig) – Training configuration
diffusion (D3PM) – D3PM diffusion process
accelerator (Accelerator | None, default: None) – Optional pre-configured accelerator
num_timestep_samples (int, default: 10) – Number of timesteps for Monte Carlo
elbo_mode (Literal['monte_carlo', 'full'], default: 'monte_carlo') – ELBO computation mode

compute_logprobs(model, samples, **kwargs)[source][source]

Delegate to internal trainer.

Parameters:

model (Module)
samples (Tensor)
kwargs (Any)

Return type:

train(*args, **kwargs)[source][source]: Delegate training to internal trainer.

validate(*args, **kwargs)[source][source]: Delegate validation to internal trainer.

save_checkpoint(*args, **kwargs)[source][source]: Delegate checkpoint saving.

load_checkpoint(*args, **kwargs)[source][source]: Delegate checkpoint loading.

prepare_batch(batch)[source]

Prepare batch for DPO training.

Parameters:: batch (dict[str, Any]) – Collated batch from DataLoader
Return type:: tuple[Tensor, Tensor, Tensor, dict[str, Any]]
Returns:: Tuple of (chosen, rejected, margin, extra_kwargs)

train_step(batch)[source]

Execute a single training step.

Parameters:: batch (dict[str, Any]) – Collated preference batch
Return type:: dict[str, float]
Returns:: Dictionary of metrics

Flow-matching models:

class medlatents.post_training.DiscreteFlowDPOTrainer(model, ref_model, config, path, source_distribution=None, accelerator=None, vocab_size=None, num_timestep_samples=10, likelihood_strategy='trajectory', time_eps=0.001)[source][source]

Bases: medlatents.post_training.dpo.base.BaseDPOTrainer

DPO trainer for discrete flow matching models.

For discrete flow matching with mixture paths, we compute likelihood as:: log p(x_1) ≈ -E_t[L_flow(x_1, t)]

where L_flow is the flow matching loss at timestep t. This provides a tractable surrogate for the true likelihood.

We support different estimation strategies: 1. “trajectory”: Sample timesteps and average flow loss 2. “endpoint”: Use only t=1 prediction (fast but less accurate) 3. “importance”: Importance-weighted sampling of timesteps

Example

>>> from medlatents.flow_matching import MixtureDiscreteProbPath
>>> path = MixtureDiscreteProbPath(PolynomialConvexScheduler(n=1.0))
>>> model = DiscreteDiT(vocab_size=1024, ...)
>>> ref_model = copy.deepcopy(model)
>>> trainer = DiscreteFlowDPOTrainer(model, ref_model, config, path=path)
>>> trainer.train(preference_dataset)

Parameters:

model (Module)
ref_model (Module | None)
config (PostTrainingConfig)
path (MixtureDiscreteProbPath)
source_distribution (Any | None, default: None)
accelerator (Accelerator | None, default: None)
vocab_size (int | None, default: None)
num_timestep_samples (int, default: 10)
likelihood_strategy (Literal['trajectory', 'endpoint', 'importance'], default: 'trajectory')
time_eps (float, default: 0.001)

__init__(model, ref_model, config, path, source_distribution=None, accelerator=None, vocab_size=None, num_timestep_samples=10, likelihood_strategy='trajectory', time_eps=0.001)[source][source]

Initialize discrete flow DPO trainer.

Parameters:

model (Module) – Flow matching model to train
ref_model (Module | None) – Frozen reference model
config (PostTrainingConfig) – Training configuration
path (MixtureDiscreteProbPath) – Discrete probability path (e.g., MixtureDiscreteProbPath)
source_distribution (Any | None, default: None) – Source distribution for flow (mask/uniform)
accelerator (Accelerator | None, default: None) – Optional pre-configured accelerator
vocab_size (int | None, default: None) – Vocabulary size
num_timestep_samples (int, default: 10) – Number of timesteps for trajectory estimation
likelihood_strategy (Literal['trajectory', 'endpoint', 'importance'], default: 'trajectory') – How to estimate likelihood
time_eps (float, default: 0.001) – Small epsilon to avoid t=0 or t=1 exactly

compute_logprobs(model, samples, **kwargs)[source][source]

Compute flow-based log-probability.

Uses the specified strategy to estimate log p(x_1).

Parameters:

model (Module) – Flow matching model
samples (Tensor) – Target token sequences (x_1) [batch, seq_len]
**kwargs (Any) – Additional model arguments

Return type:

Returns:

Log-probabilities [batch]

load_checkpoint(path)[source]

Load training checkpoint.

Parameters:: path (str) – Path to checkpoint file
Return type:: None

prepare_batch(batch)[source]

Prepare batch for DPO training.

Parameters:: batch (dict[str, Any]) – Collated batch from DataLoader
Return type:: tuple[Tensor, Tensor, Tensor, dict[str, Any]]
Returns:: Tuple of (chosen, rejected, margin, extra_kwargs)

save_checkpoint(name='latest')[source]

Save training checkpoint.

Parameters:: name (str, default: 'latest') – Checkpoint name (e.g., “latest”, “best”, “step_1000”)
Return type:: None

train(dataset, val_dataset=None)[source]

Run full training loop.

Parameters:

dataset (PreferenceDataset) – Training preference dataset
val_dataset (PreferenceDataset | None, default: None) – Optional validation dataset

Return type:

Returns:

Dictionary of training history

train_step(batch)[source]

Execute a single training step.

Parameters:: batch (dict[str, Any]) – Collated preference batch
Return type:: dict[str, float]
Returns:: Dictionary of metrics

validate(val_loader)[source]

Run validation.

Parameters:: val_loader (DataLoader) – Validation dataloader
Return type:: dict[str, float]
Returns:: Validation metrics

class medlatents.post_training.FlowDPOTrainer(model, ref_model, config, path=None, source_distribution=None, accelerator=None, vocab_size=None, num_timestep_samples=10, likelihood_strategy='trajectory', **kwargs)[source][source]

Bases: medlatents.post_training.dpo.base.BaseDPOTrainer

Unified DPO trainer for flow matching models.

This is a convenience class that handles discrete flow matching and can be extended for continuous flow matching (rectified flows).

Example

>>> trainer = FlowDPOTrainer(
...     model=flow_model,
...     ref_model=ref_flow,
...     config=config,
...     path=mixture_path,
... )
>>> trainer.train(preference_dataset)

Parameters:

model (Module)
ref_model (Module | None)
config (PostTrainingConfig)
path (MixtureDiscreteProbPath | None, default: None)
source_distribution (Any | None, default: None)
accelerator (Accelerator | None, default: None)
vocab_size (int | None, default: None)
num_timestep_samples (int, default: 10)
likelihood_strategy (Literal['trajectory', 'endpoint', 'importance'], default: 'trajectory')
kwargs (Any)

__init__(model, ref_model, config, path=None, source_distribution=None, accelerator=None, vocab_size=None, num_timestep_samples=10, likelihood_strategy='trajectory', **kwargs)[source][source]

Initialize flow DPO trainer.

Parameters:

model (Module) – Flow matching model to train
ref_model (Module | None) – Frozen reference model
config (PostTrainingConfig) – Training configuration
path (MixtureDiscreteProbPath | None, default: None) – Discrete probability path (creates default if None)
source_distribution (Any | None, default: None) – Source distribution
accelerator (Accelerator | None, default: None) – Optional pre-configured accelerator
vocab_size (int | None, default: None) – Vocabulary size
num_timestep_samples (int, default: 10) – Number of timesteps for estimation
likelihood_strategy (Literal['trajectory', 'endpoint', 'importance'], default: 'trajectory') – Likelihood estimation strategy
**kwargs (Any) – Additional arguments

compute_logprobs(model, samples, **kwargs)[source][source]

Delegate to internal trainer.

Parameters:

model (Module)
samples (Tensor)
kwargs (Any)

Return type:

train(*args, **kwargs)[source][source]: Delegate training to internal trainer.

validate(*args, **kwargs)[source][source]: Delegate validation to internal trainer.

save_checkpoint(*args, **kwargs)[source][source]: Delegate checkpoint saving.

load_checkpoint(*args, **kwargs)[source][source]: Delegate checkpoint loading.

prepare_batch(batch)[source]

Prepare batch for DPO training.

Parameters:: batch (dict[str, Any]) – Collated batch from DataLoader
Return type:: tuple[Tensor, Tensor, Tensor, dict[str, Any]]
Returns:: Tuple of (chosen, rejected, margin, extra_kwargs)

train_step(batch)[source]

Execute a single training step.

Parameters:: batch (dict[str, Any]) – Collated preference batch
Return type:: dict[str, float]
Returns:: Dictionary of metrics

Step-by-step Preference Optimization (SPO):

class medlatents.post_training.SPOTrainer(model, config, diffusion=None, flow_path=None, step_scorer=None, accelerator=None, vocab_size=None, num_candidates=4, num_steps=50)[source][source]

Bases: object

Step-by-step Preference Optimization trainer.

SPO operates by: 1. At each denoising step, generate K candidate states 2. Score candidates with a step-aware preference model 3. Apply DPO loss using best vs worst candidates 4. Randomly select one candidate to continue generation

This provides fine-grained supervision throughout the generation process, leading to better alignment than end-to-end DPO.

Example

>>> # For D3PM diffusion
>>> trainer = SPOTrainer(
...     model=backbone_model,
...     config=config,
...     diffusion=d3pm,
...     step_scorer=step_preference_model,
... )
>>> trainer.train(prompt_dataset)

Parameters:

model (Module)
config (PostTrainingConfig)
diffusion (Any | None, default: None)
flow_path (Any | None, default: None)
step_scorer (Module | None, default: None)
accelerator (Accelerator | None, default: None)
vocab_size (int | None, default: None)
num_candidates (int, default: 4)
num_steps (int, default: 50)

__init__(model, config, diffusion=None, flow_path=None, step_scorer=None, accelerator=None, vocab_size=None, num_candidates=4, num_steps=50)[source][source]

Initialize SPO trainer.

Parameters:

model (Module) – Generative model to train
config (PostTrainingConfig) – Training configuration
diffusion (Any | None, default: None) – D3PM diffusion process (for diffusion models)
flow_path (Any | None, default: None) – Flow matching path (for flow models)
step_scorer (Module | None, default: None) – Model to score intermediate states
accelerator (Accelerator | None, default: None) – Optional pre-configured accelerator
vocab_size (int | None, default: None) – Vocabulary size
num_candidates (int, default: 4) – Number of candidates per step (K)
num_steps (int, default: 50) – Number of generation steps to optimize

train_step(seq_length)[source][source]

Execute one SPO training step.

Generates a full trajectory with step-wise DPO optimization.

Parameters:: seq_length (int) – Sequence length for generation
Return type:: dict[str, float]
Returns:: Dictionary of metrics

train(seq_length=256)[source][source]

Run full SPO training loop.

Parameters:: seq_length (int, default: 256) – Sequence length for generation
Return type:: dict[str, list[float]]
Returns:: Training history

save_checkpoint(name='latest')[source][source]

Save training checkpoint.

Parameters:: name (str, default: 'latest')
Return type:: None

load_checkpoint(path)[source][source]

Load training checkpoint.

Parameters:: path (str)
Return type:: None

class medlatents.post_training.StepPreferenceModel(hidden_size=512, depth=4, num_heads=8, vocab_size=None, seq_length=256)[source][source]

Timestep-aware preference model for SPO.

This model scores intermediate states at each denoising step, allowing the model to learn step-specific preferences.

The architecture adds timestep conditioning to score states appropriately at different noise levels.

Parameters:

hidden_size (int, default: 512)
depth (int, default: 4)
num_heads (int, default: 8)
vocab_size (int | None, default: None)
seq_length (int, default: 256)

__init__(hidden_size=512, depth=4, num_heads=8, vocab_size=None, seq_length=256)[source][source]

Initialize step preference model.

Parameters:

hidden_size (int, default: 512) – Hidden dimension
depth (int, default: 4) – Number of transformer layers
num_heads (int, default: 8) – Number of attention heads
vocab_size (int | None, default: None) – Vocabulary size for discrete tokens
seq_length (int, default: 256) – Maximum sequence length

forward(x, t)[source][source]

Score a state at timestep t.

Parameters:

x (Tensor) – Token sequence [batch, seq_len]
t (Tensor) – Timestep [batch]

Return type:

Returns:

Scores [batch, 1]

T_destination = ~T_destination

add_module(name, module)[source]

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

name (str) – name of the child module. The child module can be accessed from this module using the given name
module (Module) – child module to be added to the module.

Return type:

apply(fn)[source]

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Parameters:: fn (Module -> None) – function to be applied to each submodule
Returns:: self
Return type:: Module

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

bfloat16()[source]

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

buffers(recurse=True)[source]

Return an iterator over module buffers.

Parameters:: recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.
Yields:: torch.Tensor – module buffer
Return type:: Iterator[Tensor]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

call_super_init: bool = False

children()[source]

Return an iterator over immediate children modules.

Yields:: Module – a child module
Return type:: Iterator[Module]

compile(*args, **kwargs)[source]

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()[source]

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

cuda(device=None)[source]

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

double()[source]

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

dump_patches: bool = False

eval()[source]

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:: self
Return type:: Module

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:: str

float()[source]

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

get_buffer(target)[source]

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:: target (str) – The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)
Returns:: The buffer referenced by target
Return type:: torch.Tensor
Raises:: AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:: Any extra state to store in the module’s state_dict
Return type:: object

get_parameter(target)[source]

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:: target (str) – The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)
Returns:: The Parameter referenced by target
Return type:: torch.nn.Parameter
Raises:: AttributeError – If the target string references an invalid path or resolves to something that is not an nn.Parameter

get_submodule(target)[source]

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:: target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)
Returns:: The submodule referenced by target
Return type:: torch.nn.Module
Raises:: AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half()[source]

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: self
Return type:: Module

ipu(device=None)[source]

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

load_state_dict(state_dict, strict=True, assign=False)[source]

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

state_dict (dict) – a dict containing parameters and persistent buffers.
strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

modules(remove_duplicate=True)[source]

Return an iterator over all modules in the network.

Parameters:: remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not.
Yields:: Module – a module in the network
Return type:: Iterator[Module]

Note

Duplicate modules are returned only once by default. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)[source]

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

named_buffers(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

prefix (str) – prefix to prepend to all buffer names.
recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.
remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor) – Tuple containing the name and buffer

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

named_children()[source]

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:: (str, Module) – Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

Return type:: Iterator[tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)[source]

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

memo (set[Module] | None, default: None) – a memo to store the set of modules already added to the result
prefix (str, default: '') – a prefix that will be added to the name of the module
remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

prefix (str) – prefix to prepend to all parameter names.
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

parameters(recurse=True)[source]

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:: recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
Yields:: Parameter – module parameter
Return type:: Iterator[Parameter]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

register_backward_hook(hook)[source]

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:: a handle that can be used to remove the added hook by calling handle.remove()
Return type:: torch.utils.hooks.RemovableHandle
Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], tuple[Tensor, ...] | Tensor | None])

register_buffer(name, tensor, persistent=True)[source]

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

name (str) – name of the buffer. The buffer can be accessed from this module using the given name
tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool) – whether the buffer is part of this module’s state_dict.

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)[source]

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Parameters:

hook (Callable) – The user defined hook to be registered.
prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False
with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False
always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)[source]

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Parameters:

hook (Callable) – The user defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False
with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)[source]

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)[source]

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:: a handle that can be used to remove the added hook by calling handle.remove()
Return type:: torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters:: hook (Callable) – Callable hook that will be invoked before loading the state dict.

register_module(name, module)[source]

Alias for add_module().

Parameters:

name (str)
module (Optional[Module])

Return type:

register_parameter(name, param)[source]

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

name (str) – name of the parameter. The parameter can be accessed from this module using the given name
param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Return type:

register_state_dict_post_hook(hook)[source]

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)[source]

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)[source]

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:: requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.
Returns:: self
Return type:: Module

set_extra_state(state)[source]

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:: state (dict) – Extra state from the state_dict
Return type:: None

set_submodule(target, module, strict=False)[source]

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Parameters:

target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)
module (Module) – The module to set the submodule to.
strict (bool, default: False) – If False, the method will replace an existing submodule or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError – If the target string is empty or if module is not an instance of nn.Module.
AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Return type:

share_memory()[source]

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)[source]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)[source]

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)[source]

to(dtype, non_blocking=False)[source]

to(tensor, non_blocking=False)[source]

to(memory_format=torch.channels_last)[source]

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

device (torch.device) – the desired device of the parameters and buffers in this module
dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module
tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)[source]

Move the parameters and buffers to the specified device without copying storage.

Parameters:

device (torch.device) – The desired device of the parameters and buffers in this module.
recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module

train(mode=True)[source]

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:: mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.
Returns:: self
Return type:: Module

type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Parameters:: dst_type (type or string) – the desired type
Returns:: self
Return type:: Module

xpu(device=None)[source]

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters:: device (int, optional) – if specified, all parameters will be copied to that device
Returns:: self
Return type:: Module

zero_grad(set_to_none=True)[source]

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Parameters:: set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.
Return type:: None

training: bool

class medlatents.post_training.SPOStepOutput(step_idx, timestep, loss, chosen_state, accuracy)[source][source]

Bases: object

Output from a single SPO step.

Variables:

step_idx – Index of the denoising step
timestep – Continuous timestep value
loss – DPO loss at this step
chosen_state – Selected state to continue from
accuracy – Whether chosen had higher reward

Parameters:

step_idx (int)
timestep (float)
loss (Tensor)
chosen_state (Tensor)
accuracy (float)

step_idx: int

timestep: float

loss: torch.Tensor

chosen_state: torch.Tensor

accuracy: float

__init__(step_idx, timestep, loss, chosen_state, accuracy)[source]

Parameters:

step_idx (int)
timestep (float)
loss (Tensor)
chosen_state (Tensor)
accuracy (float)

Reinforcement Learning

class medlatents.post_training.rl.BaseRLTrainer(model, reward_fn, config, ref_model=None, value_model=None, accelerator=None)[source][source]

Bases: abc.ABC

Abstract base class for RL post-training methods.

Provides common infrastructure for: - Trajectory generation and storage - Reward computation - Policy gradient estimation - Training loop with logging

Subclasses must implement: - generate_trajectories: Model-specific generation - compute_policy_loss: Method-specific loss computation

Parameters:

model (Module)
reward_fn (Callable[[Tensor], Tensor])
config (PostTrainingConfig)
ref_model (Module | None, default: None)
value_model (Module | None, default: None)
accelerator (Accelerator | None, default: None)

__init__(model, reward_fn, config, ref_model=None, value_model=None, accelerator=None)[source][source]

Initialize RL trainer.

Parameters:

model (Module) – Policy model to train
reward_fn (Callable[[Tensor], Tensor]) – Function that computes rewards for samples
config (PostTrainingConfig) – Training configuration
ref_model (Module | None, default: None) – Optional reference model for KL penalty
value_model (Module | None, default: None) – Optional value function for advantage estimation
accelerator (Accelerator | None, default: None) – Optional pre-configured accelerator

abstractmethod generate_trajectories(batch_size, seq_length, **kwargs)[source][source]

Generate trajectories by sampling from the model.

Parameters:

batch_size (int) – Number of trajectories to generate
seq_length (int) – Sequence length
**kwargs (Any) – Additional generation arguments

Return type:

Returns:

List of Trajectory objects

abstractmethod compute_policy_loss(trajectories)[source][source]

Compute policy gradient loss.

Parameters:: trajectories (list[Trajectory]) – Batch of trajectories
Return type:: tuple[Tensor, dict[str, float]]
Returns:: Tuple of (loss, metrics_dict)

compute_kl_penalty(log_probs, ref_log_probs)[source][source]

Compute KL divergence penalty.

Parameters:

log_probs (Tensor) – Log probs from current policy
ref_log_probs (Tensor) – Log probs from reference policy

Return type:

Returns:

KL divergence estimate

train_step(seq_length)[source][source]

Execute one training step.

Generate trajectories
Compute rewards
Compute advantages
Update policy

Parameters:: seq_length (int) – Sequence length for generation
Return type:: dict[str, float]
Returns:: Dictionary of metrics

train(seq_length=256)[source][source]

Run full RL training loop.

Parameters:: seq_length (int, default: 256) – Sequence length for generation
Return type:: dict[str, list[float]]
Returns:: Training history

save_checkpoint(name='latest')[source][source]

Save training checkpoint.

Parameters:: name (str, default: 'latest')
Return type:: None

load_checkpoint(path)[source][source]

Load training checkpoint.

Parameters:: path (str)
Return type:: None

class medlatents.post_training.rl.DDPOTrainer(model, reward_fn, config, diffusion, ref_model=None, value_model=None, accelerator=None, num_inference_steps=50, clip_range=None, kl_coeff=None)[source][source]

Bases: medlatents.post_training.rl.base.BaseRLTrainer

DDPO trainer for diffusion models.

Implements PPO-style policy gradient optimization for diffusion models. The key idea is to: 1. Sample full denoising trajectories 2. Compute rewards for final samples 3. Estimate advantages using GAE 4. Update policy with clipped surrogate objective

Example

>>> from medlatents.diffusion import D3PM
>>> d3pm = D3PM(num_classes=1024, num_timesteps=1000)
>>> trainer = DDPOTrainer(
...     model=backbone,
...     reward_fn=reward_model,
...     config=config,
...     diffusion=d3pm,
... )
>>> trainer.train(seq_length=256)

Parameters:

model (Module)
reward_fn (Callable[[Tensor], Tensor])
config (PostTrainingConfig)
diffusion (D3PM)
ref_model (Module | None, default: None)
value_model (Module | None, default: None)
accelerator (Accelerator | None, default: None)
num_inference_steps (int, default: 50)
clip_range (float | None, default: None)
kl_coeff (float | None, default: None)

__init__(model, reward_fn, config, diffusion, ref_model=None, value_model=None, accelerator=None, num_inference_steps=50, clip_range=None, kl_coeff=None)[source][source]

Initialize DDPO trainer.

Parameters:

model (Module) – Diffusion backbone model
reward_fn (Callable[[Tensor], Tensor]) – Reward function for samples
config (PostTrainingConfig) – Training configuration
diffusion (D3PM) – D3PM diffusion process
ref_model (Module | None, default: None) – Reference model for KL penalty
value_model (Module | None, default: None) – Value function for advantage estimation
accelerator (Accelerator | None, default: None) – Optional accelerator
num_inference_steps (int, default: 50) – Number of denoising steps
clip_range (float | None, default: None) – PPO clipping range (uses config if None)
kl_coeff (float | None, default: None) – KL penalty coefficient (uses config if None)

generate_trajectories(batch_size, seq_length, **kwargs)[source][source]

Generate denoising trajectories.

Samples full denoising paths while recording states, actions, and log probabilities at each step.

Parameters:

batch_size (int) – Number of trajectories
seq_length (int) – Sequence length
**kwargs (Any) – Additional arguments

Return type:

Returns:

List of Trajectory objects

compute_policy_loss(trajectories)[source][source]

Compute PPO-style clipped surrogate loss.

The PPO objective is:: L = -min(r*A, clip(r, 1-ε, 1+ε)*A)

where r = π(a|s) / π_old(a|s) is the probability ratio.

Parameters:: trajectories (list[Trajectory]) – List of trajectories with advantages
Return type:: tuple[Tensor, dict[str, float]]
Returns:: Tuple of (loss, metrics)

compute_kl_penalty(log_probs, ref_log_probs)[source]

Compute KL divergence penalty.

Parameters:

log_probs (Tensor) – Log probs from current policy
ref_log_probs (Tensor) – Log probs from reference policy

Return type:

Returns:

KL divergence estimate

load_checkpoint(path)[source]

Load training checkpoint.

Parameters:: path (str)
Return type:: None

save_checkpoint(name='latest')[source]

Save training checkpoint.

Parameters:: name (str, default: 'latest')
Return type:: None

train(seq_length=256)[source]

Run full RL training loop.

Parameters:: seq_length (int, default: 256) – Sequence length for generation
Return type:: dict[str, list[float]]
Returns:: Training history

train_step(seq_length)[source]

Execute one training step.

Generate trajectories
Compute rewards
Compute advantages
Update policy

Parameters:: seq_length (int) – Sequence length for generation
Return type:: dict[str, float]
Returns:: Dictionary of metrics

class medlatents.post_training.rl.DDPODiscreteFlowTrainer(model, reward_fn, config, path, source_distribution=None, ref_model=None, accelerator=None, vocab_size=None, num_inference_steps=50)[source][source]

Bases: medlatents.post_training.rl.base.BaseRLTrainer

DDPO trainer for discrete flow matching models.

Adapts DDPO for discrete flow matching where the model predicts token distributions at each timestep along the flow.

Parameters:

model (Module)
reward_fn (Callable[[Tensor], Tensor])
config (PostTrainingConfig)
path (Any)
source_distribution (Any | None, default: None)
ref_model (Module | None, default: None)
accelerator (Accelerator | None, default: None)
vocab_size (int | None, default: None)
num_inference_steps (int, default: 50)

__init__(model, reward_fn, config, path, source_distribution=None, ref_model=None, accelerator=None, vocab_size=None, num_inference_steps=50)[source][source]

Initialize discrete flow DDPO trainer.

Parameters:

model (Module)
reward_fn (Callable[[Tensor], Tensor])
config (PostTrainingConfig)
path (Any)
source_distribution (Any | None, default: None)
ref_model (Module | None, default: None)
accelerator (Accelerator | None, default: None)
vocab_size (int | None, default: None)
num_inference_steps (int, default: 50)

generate_trajectories(batch_size, seq_length, **kwargs)[source][source]

Generate flow trajectories.

Parameters:

batch_size (int)
seq_length (int)
kwargs (Any)

Return type:

compute_policy_loss(trajectories)[source][source]

Compute policy loss (same as DDPOTrainer).

Parameters:: trajectories (list[Trajectory])
Return type:: tuple[Tensor, dict[str, float]]

compute_kl_penalty(log_probs, ref_log_probs)[source]

Compute KL divergence penalty.

Parameters:

log_probs (Tensor) – Log probs from current policy
ref_log_probs (Tensor) – Log probs from reference policy

Return type:

Returns:

KL divergence estimate

load_checkpoint(path)[source]

Load training checkpoint.

Parameters:: path (str)
Return type:: None

save_checkpoint(name='latest')[source]

Save training checkpoint.

Parameters:: name (str, default: 'latest')
Return type:: None

train(seq_length=256)[source]

Run full RL training loop.

Parameters:: seq_length (int, default: 256) – Sequence length for generation
Return type:: dict[str, list[float]]
Returns:: Training history

train_step(seq_length)[source]

Execute one training step.

Generate trajectories
Compute rewards
Compute advantages
Update policy

Parameters:: seq_length (int) – Sequence length for generation
Return type:: dict[str, float]
Returns:: Dictionary of metrics

class medlatents.post_training.rl.GRPOTrainer(model, reward_fn, config, ref_model=None, accelerator=None, num_samples_per_prompt=None, vocab_size=None, num_inference_steps=50, model_type='diffusion', diffusion=None, flow_path=None)[source][source]

Bases: medlatents.post_training.rl.base.BaseRLTrainer

Group Relative Policy Optimization trainer.

GRPO samples multiple outputs for each input and computes advantages relative to the group mean. This eliminates the need for a learned value function while providing low-variance gradients.

The GRPO objective is:: L = -E[A(x) * log π(x)] + β * KL(π || π_ref)

where A(x) = (r(x) - mean(r)) / std(r) is the group-relative advantage.

Example

>>> trainer = GRPOTrainer(
...     model=model,
...     reward_fn=reward_model,
...     config=config,
...     num_samples_per_prompt=8,
... )
>>> trainer.train(seq_length=256)

Parameters:

model (Module)
reward_fn (Callable[[Tensor], Tensor])
config (PostTrainingConfig)
ref_model (Module | None, default: None)
accelerator (Accelerator | None, default: None)
num_samples_per_prompt (int | None, default: None)
vocab_size (int | None, default: None)
num_inference_steps (int, default: 50)
model_type (str, default: 'diffusion')
diffusion (Any | None, default: None)
flow_path (Any | None, default: None)

__init__(model, reward_fn, config, ref_model=None, accelerator=None, num_samples_per_prompt=None, vocab_size=None, num_inference_steps=50, model_type='diffusion', diffusion=None, flow_path=None)[source][source]

Initialize GRPO trainer.

Parameters:

model (Module) – Model to train
reward_fn (Callable[[Tensor], Tensor]) – Reward function
config (PostTrainingConfig) – Training configuration
ref_model (Module | None, default: None) – Reference model for KL penalty
accelerator (Accelerator | None, default: None) – Optional accelerator
num_samples_per_prompt (int | None, default: None) – Samples per group (uses config if None)
vocab_size (int | None, default: None) – Vocabulary size
num_inference_steps (int, default: 50) – Number of generation steps
model_type (str, default: 'diffusion') – “diffusion”, “flow”, or “maskgit”
diffusion (Any | None, default: None) – D3PM for diffusion models
flow_path (Any | None, default: None) – Path for flow models

generate_trajectories(batch_size, seq_length, **kwargs)[source][source]

Generate trajectories for GRPO.

Generates multiple samples per “prompt” for group-relative advantages.

Parameters:

batch_size (int)
seq_length (int)
kwargs (Any)

Return type:

compute_policy_loss(trajectories)[source][source]

Compute GRPO loss with group-relative advantages.

Groups trajectories by prompt and computes advantages relative to the group mean reward.

Parameters:: trajectories (list[Trajectory])
Return type:: tuple[Tensor, dict[str, float]]

compute_kl_penalty(log_probs, ref_log_probs)[source]

Compute KL divergence penalty.

Parameters:

log_probs (Tensor) – Log probs from current policy
ref_log_probs (Tensor) – Log probs from reference policy

Return type:

Returns:

KL divergence estimate

load_checkpoint(path)[source]

Load training checkpoint.

Parameters:: path (str)
Return type:: None

save_checkpoint(name='latest')[source]

Save training checkpoint.

Parameters:: name (str, default: 'latest')
Return type:: None

train(seq_length=256)[source]

Run full RL training loop.

Parameters:: seq_length (int, default: 256) – Sequence length for generation
Return type:: dict[str, list[float]]
Returns:: Training history

train_step(seq_length)[source]

Execute one training step.

Generate trajectories
Compute rewards
Compute advantages
Update policy

Parameters:: seq_length (int) – Sequence length for generation
Return type:: dict[str, float]
Returns:: Dictionary of metrics

class medlatents.post_training.rl.GARDOTrainer(model, reward_fn, config, ref_model=None, accelerator=None, reward_threshold=None, distance_penalty=None, vocab_size=None, num_inference_steps=50, model_type='diffusion', diffusion=None, adaptive_kl=True, kl_target=0.1)[source][source]

Bases: medlatents.post_training.rl.base.BaseRLTrainer

Gradient-Aware Reward Distance Optimization trainer.

GARDO improves upon standard reward fine-tuning by: 1. Monitoring gradient alignment between reward and KL terms 2. Adaptively scaling the KL coefficient 3. Filtering updates based on reward threshold

The objective is:: L = -r(x) + α(grad) * KL(π || π_ref)

where α(grad) is dynamically adjusted based on gradient cosine similarity.

Example

>>> trainer = GARDOTrainer(
...     model=model,
...     reward_fn=reward_model,
...     config=config,
...     reward_threshold=0.1,
... )
>>> trainer.train(seq_length=256)

Parameters:

model (Module)
reward_fn (Callable[[Tensor], Tensor])
config (PostTrainingConfig)
ref_model (Module | None, default: None)
accelerator (Accelerator | None, default: None)
reward_threshold (float | None, default: None)
distance_penalty (float | None, default: None)
vocab_size (int | None, default: None)
num_inference_steps (int, default: 50)
model_type (str, default: 'diffusion')
diffusion (Any | None, default: None)
adaptive_kl (bool, default: True)
kl_target (float, default: 0.1)

__init__(model, reward_fn, config, ref_model=None, accelerator=None, reward_threshold=None, distance_penalty=None, vocab_size=None, num_inference_steps=50, model_type='diffusion', diffusion=None, adaptive_kl=True, kl_target=0.1)[source][source]

Initialize GARDO trainer.

Parameters:

model (Module) – Model to train
reward_fn (Callable[[Tensor], Tensor]) – Reward function
config (PostTrainingConfig) – Training configuration
ref_model (Module | None, default: None) – Reference model for KL computation
accelerator (Accelerator | None, default: None) – Optional accelerator
reward_threshold (float | None, default: None) – Minimum reward improvement to update
distance_penalty (float | None, default: None) – Base KL penalty coefficient
vocab_size (int | None, default: None) – Vocabulary size
num_inference_steps (int, default: 50) – Generation steps
model_type (str, default: 'diffusion') – “diffusion”, “flow”, or “maskgit”
diffusion (Any | None, default: None) – D3PM for diffusion models
adaptive_kl (bool, default: True) – Whether to adapt KL coefficient
kl_target (float, default: 0.1) – Target KL divergence for adaptive scaling

generate_trajectories(batch_size, seq_length, **kwargs)[source][source]

Generate trajectories.

Parameters:

batch_size (int)
seq_length (int)
kwargs (Any)

Return type:

compute_policy_loss(trajectories)[source][source]

Compute GARDO loss with adaptive KL penalty.

The GARDO loss combines: 1. Negative reward (to maximize reward) 2. KL penalty (to stay close to reference) 3. Reward thresholding (to focus on improvements)

Parameters:: trajectories (list[Trajectory])
Return type:: tuple[Tensor, dict[str, float]]

compute_kl_penalty(log_probs, ref_log_probs)[source]

Compute KL divergence penalty.

Parameters:

log_probs (Tensor) – Log probs from current policy
ref_log_probs (Tensor) – Log probs from reference policy

Return type:

Returns:

KL divergence estimate

load_checkpoint(path)[source]

Load training checkpoint.

Parameters:: path (str)
Return type:: None

save_checkpoint(name='latest')[source]

Save training checkpoint.

Parameters:: name (str, default: 'latest')
Return type:: None

train(seq_length=256)[source]

Run full RL training loop.

Parameters:: seq_length (int, default: 256) – Sequence length for generation
Return type:: dict[str, list[float]]
Returns:: Training history

train_step(seq_length)[source]

Execute one training step.

Generate trajectories
Compute rewards
Compute advantages
Update policy

Parameters:: seq_length (int) – Sequence length for generation
Return type:: dict[str, float]
Returns:: Dictionary of metrics

class medlatents.post_training.rl.Trajectory(states, actions, timesteps, log_probs, rewards, final_sample, prompt=None, advantages=None, returns=None)[source][source]

Bases: object

A single generation trajectory with rewards.

For diffusion/flow models, this represents the full denoising path. For autoregressive models, this is the sequence of token generations.

Variables:

states – Intermediate states [num_steps, seq_len] or [num_steps, seq_len, dim]
actions – Actions taken (tokens generated) [num_steps, seq_len]
timesteps – Timestep at each step [num_steps]
log_probs – Log probabilities of actions [num_steps]
rewards – Rewards at each step [num_steps] (often only final is non-zero)
final_sample – The final generated sample
prompt – Optional conditioning information

Parameters:

states (Tensor)
actions (Tensor)
timesteps (Tensor)
log_probs (Tensor)
rewards (Tensor)
final_sample (Tensor)
prompt (Any, default: None)
advantages (Tensor | None, default: None)
returns (Tensor | None, default: None)

states: torch.Tensor

actions: torch.Tensor

timesteps: torch.Tensor

log_probs: torch.Tensor

rewards: torch.Tensor

final_sample: torch.Tensor

prompt: Any = None

advantages: torch.Tensor | None = None

returns: torch.Tensor | None = None

to(device)[source][source]

Move trajectory to device.

Parameters:: device (device)
Return type:: Trajectory

__init__(states, actions, timesteps, log_probs, rewards, final_sample, prompt=None, advantages=None, returns=None)[source]

Parameters:

states (Tensor)
actions (Tensor)
timesteps (Tensor)
log_probs (Tensor)
rewards (Tensor)
final_sample (Tensor)
prompt (Any, default: None)
advantages (Tensor | None, default: None)
returns (Tensor | None, default: None)

class medlatents.post_training.rl.TrajectoryBuffer(max_size=10000, device=device(type='cpu'))[source][source]

Bases: object

Buffer for storing and sampling trajectories.

Implements experience replay for RL training with support for priority sampling and trajectory filtering.

Parameters:

max_size (int, default: 10000)
device (device, default: device(type='cpu'))

__init__(max_size=10000, device=device(type='cpu'))[source][source]

Initialize trajectory buffer.

Parameters:

max_size (int, default: 10000) – Maximum number of trajectories to store
device (device, default: device(type='cpu')) – Device for storing trajectories

buffer: list[medlatents.post_training.rl.base.Trajectory]

priorities: list[float]

add(trajectory, priority=1.0)[source][source]

Add a trajectory to the buffer.

Parameters:

trajectory (Trajectory)
priority (float, default: 1.0)

Return type:

sample(batch_size, prioritized=False)[source][source]

Sample a batch of trajectories.

Parameters:

batch_size (int) – Number of trajectories to sample
prioritized (bool, default: False) – If True, sample proportional to priority

Return type:

Returns:

List of sampled trajectories

clear()[source][source]

Clear the buffer.

Return type:: None

Distillation

class medlatents.post_training.distillation.ReflowTrainer(model, config, path, vocab_size, accelerator=None, num_inference_steps=50)[source][source]

Bases: object

Iterative reflow training for flow matching models.

Implements the reflow procedure: 1. Generate pairs from current model 2. Train on pairs to get straighter trajectories 3. Repeat

Each iteration produces more linear trajectories, enabling faster sampling with fewer steps.

Example

>>> trainer = ReflowTrainer(
...     model=flow_model,
...     config=config,
...     path=flow_path,
...     vocab_size=1024,
... )
>>> # Run one reflow iteration
>>> trainer.reflow_iteration(data_loader)
>>>
>>> # Run multiple iterations
>>> for i in range(3):
...     trainer.reflow_iteration(data_loader)

Parameters:

model (Module)
config (PostTrainingConfig)
path (Any)
vocab_size (int)
accelerator (Accelerator | None, default: None)
num_inference_steps (int, default: 50)

__init__(model, config, path, vocab_size, accelerator=None, num_inference_steps=50)[source][source]

Initialize reflow trainer.

Parameters:

model (Module) – Flow model to train
config (PostTrainingConfig) – Training configuration
path (Any) – Flow probability path
vocab_size (int) – Vocabulary size
accelerator (Accelerator | None, default: None) – Optional accelerator
num_inference_steps (int, default: 50) – Inference steps for pair generation

train_on_pairs(noise_samples, data_samples, num_steps=None)[source][source]

Train on generated (noise, data) pairs.

Parameters:

noise_samples (Tensor) – Source noise [N, seq_len]
data_samples (Tensor) – Target data [N, seq_len]
num_steps (int | None, default: None) – Training steps (uses config if None)

Return type:

Returns:

Training history

reflow_iteration(data_loader, max_pairs=None, num_train_steps=None)[source][source]

Run one complete reflow iteration.

Collect data samples
Generate (noise, data) pairs
Train on pairs

Parameters:

data_loader (DataLoader) – DataLoader with training data
max_pairs (int | None, default: None) – Maximum pairs to generate
num_train_steps (int | None, default: None) – Training steps for this iteration

Return type:

dict[str, Any]

Returns:

Dictionary with iteration metrics

train(data_loader, num_iterations=None)[source][source]

Run full reflow training with multiple iterations.

Parameters:

data_loader (DataLoader) – DataLoader with training data
num_iterations (int | None, default: None) – Number of reflow iterations

Return type:

dict[str, list[Any]]

Returns:

Training history across all iterations

save_checkpoint(name='latest')[source][source]

Save checkpoint.

Parameters:: name (str, default: 'latest')
Return type:: None

load_checkpoint(path)[source][source]

Load checkpoint.

Parameters:: path (str)
Return type:: None

class medlatents.post_training.distillation.ReflowPairGenerator(model, path, vocab_size, device, num_inference_steps=50)[source][source]

Bases: object

Generate (noise, data) pairs for reflow training.

Uses the current flow model to generate aligned pairs that represent straight-line trajectories in probability space.

Parameters:

model (Module)
path (Any)
vocab_size (int)
device (device)
num_inference_steps (int, default: 50)

__init__(model, path, vocab_size, device, num_inference_steps=50)[source][source]

Initialize pair generator.

Parameters:

model (Module) – Current flow model
path (Any) – Flow probability path
vocab_size (int) – Vocabulary size for discrete models
device (device) – Computation device
num_inference_steps (int, default: 50) – Steps for generation

generate_pairs(data_samples, batch_size=32)[source][source]

Generate (noise, data) pairs from real data.

For each data sample x_1: 1. Sample noise x_0 from source distribution 2. Generate trajectory: x_0 -> x_1’ (using model) 3. Store pair (x_0, x_1)

The key insight is that we use the REAL x_1 as target, not the model’s generated x_1’, creating straighter paths.

Parameters:

data_samples (Tensor) – Real data samples [N, seq_len]
batch_size (int, default: 32) – Batch size for generation

Return type:

tuple[Tensor, Tensor]

Returns:

Tuple of (noise_samples, data_samples) for training

class medlatents.post_training.distillation.ConsistencyTrainer(student, config, teacher=None, diffusion=None, accelerator=None, mode='distillation', sigma_min=0.002, sigma_max=80.0, sigma_data=0.5, s0=10, s1=1280, huber_c=None)[source][source]

Bases: object

Consistency training/distillation for diffusion models.

Supports two modes: 1. Consistency Training (CT): Train from scratch with self-consistency 2. Consistency Distillation (CD): Distill from a pretrained teacher

Example

>>> # Consistency Distillation
>>> trainer = ConsistencyTrainer(
...     student=student_model,
...     teacher=teacher_model,
...     config=config,
...     mode="distillation",
... )
>>> trainer.train(train_loader)
>>>
>>> # Consistency Training
>>> trainer = ConsistencyTrainer(
...     student=model,
...     config=config,
...     mode="training",
... )
>>> trainer.train(train_loader)

Parameters:

student (Module)
config (PostTrainingConfig)
teacher (Module | None, default: None)
diffusion (Any | None, default: None)
accelerator (Accelerator | None, default: None)
mode (Literal['training', 'distillation'], default: 'distillation')
sigma_min (float, default: 0.002)
sigma_max (float, default: 80.0)
sigma_data (float, default: 0.5)
s0 (int, default: 10)
s1 (int, default: 1280)
huber_c (float | None, default: None)

__init__(student, config, teacher=None, diffusion=None, accelerator=None, mode='distillation', sigma_min=0.002, sigma_max=80.0, sigma_data=0.5, s0=10, s1=1280, huber_c=None)[source][source]

Initialize consistency trainer.

Parameters:

student (Module) – Model to train
config (PostTrainingConfig) – Training configuration
teacher (Module | None, default: None) – Teacher model for distillation (required if mode=”distillation”)
diffusion (Any | None, default: None) – D3PM for discrete models
accelerator (Accelerator | None, default: None) – Optional accelerator
mode (Literal['training', 'distillation'], default: 'distillation') – “training” or “distillation”
sigma_min (float, default: 0.002) – Minimum noise level
sigma_max (float, default: 80.0) – Maximum noise level
sigma_data (float, default: 0.5) – Data standard deviation
s0 (int, default: 10) – Initial discretization steps
s1 (int, default: 1280) – Final discretization steps
huber_c (float | None, default: None) – Pseudo-Huber constant

train_step_distillation(batch)[source][source]

One step of consistency distillation.

Uses teacher to generate targets for student.

Parameters:: batch (Tensor)
Return type:: dict[str, float]

train_step_training(batch)[source][source]

One step of consistency training (no teacher).

Uses self-consistency constraint.

Parameters:: batch (Tensor)
Return type:: dict[str, float]

train(train_loader)[source][source]

Run consistency training loop.

Parameters:: train_loader (Any) – DataLoader for training data
Return type:: dict[str, list[float]]
Returns:: Training history

save_checkpoint(name='latest')[source][source]

Save checkpoint.

Parameters:: name (str, default: 'latest')
Return type:: None

load_checkpoint(path)[source][source]

Load checkpoint.

Parameters:: path (str)
Return type:: None

medlatents.post_training.distillation.pseudo_huber_loss(x, y, c=0.00054)[source][source]

Pseudo-Huber loss for consistency training.

More robust than MSE for high-dimensional data. L(x, y) = sqrt((x - y)^2 + c^2) - c

Parameters:

x (Tensor) – Predictions
y (Tensor) – Targets
c (float, default: 0.00054) – Huber constant (controls transition from L2 to L1)

Return type:

Returns:

Loss value

Self-Play

class medlatents.post_training.self_play.SPINTrainer(model, config, accelerator=None, vocab_size=None, generate_fn=None, num_generation_steps=50, model_type='maskgit')[source][source]

Bases: object

Self-Play Iterative Training (SPIN) for generative models.

SPIN improves a model by creating a curriculum of increasingly difficult discrimination tasks:

Iteration 1: Distinguish ground truth from weak model outputs Iteration 2: Distinguish ground truth from stronger model outputs …

This is implemented as iterative DPO where: - Chosen = ground truth data - Rejected = model’s own generations

Example

>>> trainer = SPINTrainer(
...     model=model,
...     config=config,
...     vocab_size=1024,
... )
>>> trainer.train(data_loader, num_iterations=3)

Parameters:

model (Module)
config (PostTrainingConfig)
accelerator (Accelerator | None, default: None)
vocab_size (int | None, default: None)
generate_fn (Callable[[Module, int, int], Tensor] | None, default: None)
num_generation_steps (int, default: 50)
model_type (str, default: 'maskgit')

__init__(model, config, accelerator=None, vocab_size=None, generate_fn=None, num_generation_steps=50, model_type='maskgit')[source][source]

Initialize SPIN trainer.

Parameters:

model (Module) – Model to train
config (PostTrainingConfig) – Training configuration
accelerator (Accelerator | None, default: None) – Optional accelerator
vocab_size (int | None, default: None) – Vocabulary size
generate_fn (Callable[[Module, int, int], Tensor] | None, default: None) – Custom generation function (model, batch_size, seq_len) -> samples
num_generation_steps (int, default: 50) – Steps for generation
model_type (str, default: 'maskgit') – Model type for default generation

train_step(real_data)[source][source]

One SPIN training step.

Parameters:: real_data (Tensor) – Ground truth samples (chosen)
Return type:: dict[str, float]
Returns:: Metrics dictionary

train_iteration(data_loader, num_steps=None)[source][source]

Run one SPIN iteration.

Parameters:

data_loader (DataLoader) – DataLoader with ground truth data
num_steps (int | None, default: None) – Steps for this iteration

Return type:

Returns:

Training history

train(data_loader, num_iterations=None)[source][source]

Run full SPIN training with multiple iterations.

Parameters:

data_loader (DataLoader) – DataLoader with ground truth data
num_iterations (int | None, default: None) – Number of self-play iterations

Return type:

dict[str, Any]

Returns:

Complete training history

save_checkpoint(name='latest')[source][source]

Save checkpoint.

Parameters:: name (str, default: 'latest')
Return type:: None

load_checkpoint(path)[source][source]

Load checkpoint.

Parameters:: path (str)
Return type:: None

class medlatents.post_training.self_play.RFTTrainer(model, reward_fn, config, accelerator=None, vocab_size=None, num_samples_per_prompt=8, top_k=None, generate_fn=None, num_generation_steps=50, model_type='maskgit')[source][source]

Bases: object

Rejection Fine-Tuning trainer.

RFT is a simple but effective method that: 1. Generates multiple samples from the model 2. Selects the best ones using a reward function 3. Fine-tunes on the selected samples

This provides quality improvement without the complexity of RL.

Example

>>> trainer = RFTTrainer(
...     model=model,
...     reward_fn=reward_model,
...     config=config,
...     num_samples_per_prompt=8,
...     top_k=2,
... )
>>> trainer.train(prompt_loader)

Parameters:

model (Module)
reward_fn (Callable[[Tensor], Tensor])
config (PostTrainingConfig)
accelerator (Accelerator | None, default: None)
vocab_size (int | None, default: None)
num_samples_per_prompt (int, default: 8)
top_k (int | None, default: None)
generate_fn (Callable[[Module, int, int], Tensor] | None, default: None)
num_generation_steps (int, default: 50)
model_type (str, default: 'maskgit')

__init__(model, reward_fn, config, accelerator=None, vocab_size=None, num_samples_per_prompt=8, top_k=None, generate_fn=None, num_generation_steps=50, model_type='maskgit')[source][source]

Initialize RFT trainer.

Parameters:

model (Module) – Model to train
reward_fn (Callable[[Tensor], Tensor]) – Function that scores samples (higher = better)
config (PostTrainingConfig) – Training configuration
accelerator (Accelerator | None, default: None) – Optional accelerator
vocab_size (int | None, default: None) – Vocabulary size
num_samples_per_prompt (int, default: 8) – Number of samples to generate per prompt
top_k (int | None, default: None) – Number of top samples to keep (default: num_samples // 2)
generate_fn (Callable[[Module, int, int], Tensor] | None, default: None) – Custom generation function
num_generation_steps (int, default: 50) – Steps for generation
model_type (str, default: 'maskgit') – Model type for generation

train_step(seq_length, num_prompts)[source][source]

One RFT training step.

Generate samples
Select top-k by reward
Fine-tune on selected

Parameters:

seq_length (int) – Sequence length
num_prompts (int) – Number of parallel generations

Return type:

dict[str, float]

Returns:

Metrics dictionary

train(seq_length=256, num_prompts_per_step=4)[source][source]

Run RFT training.

Parameters:

seq_length (int, default: 256) – Sequence length for generation
num_prompts_per_step (int, default: 4) – Prompts per training step

Return type: