burn_dragon_train 0.5.0

Training utilities for burn_dragon
Documentation
#[cfg(feature = "train")]
pub mod artifacts;
pub mod core;
pub mod optimizer;
pub mod run_layout;

#[cfg(feature = "train")]
pub use artifacts::VisionArtifactOutputMode;
pub use core::{
    FsdpMixedPrecisionKind, GdpoConfig, GdpoHardGate, KernelSpec, LayerStateSpec, LowBitMemorySpec,
    LowBitModelSpec, LowBitSavedActivationInventorySpec, LowBitSavedActivationTensorSpec,
    ModelSpec, OptimizerSpec, ParallelCheckpointConfig, ParallelCheckpointFormat,
    ParallelCommunicationBackend, ParallelConfig, ParallelDataConfig, ParallelFsdpConfig,
    ParallelPipelineCacheConfig, ParallelPipelineConfig, ParallelSpec, ParallelTensorConfig,
    ParallelismKind, PipelineCacheEvictionKind, PipelineCachePolicy, PipelineCommunicationKind,
    PipelinePartitionKind, PipelineScheduleKind, PipelineSharedWeightSyncKind,
    PipelineTransportDtype, SequenceKernelConfig, StateAxisSpec, StateLayout, StateTensorSpec,
    TensorParallelAxis, TensorParallelPartitionKind, VisionTeacherVariant, WgpuBackend,
    WgpuGenerationExecutor, WgpuInferenceConfig, WgpuMemoryConfig, WgpuRuntimeConfig,
    WgpuStartupAutotuneConfig, WgpuTrainingConfig,
};
pub use optimizer::{
    LearningRateScheduleConfig, MuonAdjustLrFn, MuonHybridConfig, OptimizerConfig, OptimizerKind,
    OptimizerScheduleMode,
};
pub use run_layout::RunLayoutConfig;