From 0a9341acdefd1164c0e4bc35b4a2716f7f078b48 Mon Sep 17 00:00:00 2001 From: Dan Saunders Date: Thu, 14 Aug 2025 01:53:24 -0400 Subject: [PATCH] nits --- src/axolotl/integrations/diffusion/__init__.py | 6 +----- src/axolotl/integrations/diffusion/args.py | 8 ++++---- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/src/axolotl/integrations/diffusion/__init__.py b/src/axolotl/integrations/diffusion/__init__.py index 84ae75ccb..942f8051f 100644 --- a/src/axolotl/integrations/diffusion/__init__.py +++ b/src/axolotl/integrations/diffusion/__init__.py @@ -1,8 +1,4 @@ -""" -Diffusion LM training plugin for Axolotl. - -This plugin enables diffusion language model training using the LLaDA approach. -""" +"""Diffusion LM training plugin init.""" from .args import DiffusionArgs from .plugin import DiffusionPlugin diff --git a/src/axolotl/integrations/diffusion/args.py b/src/axolotl/integrations/diffusion/args.py index 949b08824..639c85055 100644 --- a/src/axolotl/integrations/diffusion/args.py +++ b/src/axolotl/integrations/diffusion/args.py @@ -1,4 +1,4 @@ -"""Configuration arguments for diffusion LM training.""" +"""Config args for diffusion LM training.""" from typing import Literal @@ -8,7 +8,7 @@ from pydantic import BaseModel, Field class DiffusionArgs(BaseModel): """Arguments for diffusion LM training plugin.""" - # Noise schedule configuration + # Noise schedule config noise_schedule: Literal["linear", "cosine"] = Field( default="linear", description="Type of noise schedule for diffusion training" ) @@ -28,7 +28,7 @@ class DiffusionArgs(BaseModel): default=1000, ge=1, description="Number of diffusion timesteps" ) - # Forward process parameters + # Forward process config eps: float = Field( default=1e-3, ge=0.0, @@ -36,7 +36,7 @@ class DiffusionArgs(BaseModel): description="Epsilon value for minimum masking probability in forward process", ) - # Training configuration + # Training config importance_weighting: bool = Field( default=True, description="Apply importance weighting to loss based on masking probability",