integrations.diffusion.plugin

integrations.diffusion.plugin

Diffusion LM training plugin for Axolotl.

Classes

Name	Description
DiffusionPlugin	Plugin for diffusion language model training.

DiffusionPlugin

integrations.diffusion.plugin.DiffusionPlugin()

Plugin for diffusion language model training.

This plugin enables diffusion-based training using the LLaDA approach, which uses random masking and bidirectional attention to train language models.

Methods

Name	Description
get_input_args	Returns the pydantic model for LLaDA plugin arguments.
get_trainer_cls	Return custom trainer class for diffusion training.
post_model_load	Perform actions after model is loaded.
post_trainer_create	Configure trainer after creation.

get_input_args

integrations.diffusion.plugin.DiffusionPlugin.get_input_args()

Returns the pydantic model for LLaDA plugin arguments.

get_trainer_cls

integrations.diffusion.plugin.DiffusionPlugin.get_trainer_cls(cfg)

Return custom trainer class for diffusion training.

post_model_load

integrations.diffusion.plugin.DiffusionPlugin.post_model_load(cfg, model)

Perform actions after model is loaded.

post_trainer_create

integrations.diffusion.plugin.DiffusionPlugin.post_trainer_create(cfg, trainer)

Configure trainer after creation.