fastvideo.v1.training.training_pipeline

fastvideo.v1.training.training_pipeline#

Module Contents#

Classes#

TrainingPipeline

A pipeline for training a model. All training pipelines should inherit from this class. All reusable components and code should be implemented in this class.

Data#

GRADIENT_CHECK_DTYPE

logger

API#

fastvideo.v1.training.training_pipeline.GRADIENT_CHECK_DTYPE[source]#

None

class fastvideo.v1.training.training_pipeline.TrainingPipeline(model_path: str, fastvideo_args: Union[fastvideo.v1.fastvideo_args.FastVideoArgs, fastvideo.v1.fastvideo_args.TrainingArgs], required_config_modules: Optional[List[str]] = None, loaded_modules: Optional[Dict[str, torch.nn.Module]] = None)[source]#

Bases: fastvideo.v1.pipelines.ComposedPipelineBase, abc.ABC

A pipeline for training a model. All training pipelines should inherit from this class. All reusable components and code should be implemented in this class.

Initialization

Initialize the pipeline. After init, the pipeline should be ready to use. The pipeline should be stateless and not hold any batch state.

create_pipeline_stages(fastvideo_args: fastvideo.v1.fastvideo_args.FastVideoArgs)[source]#
current_epoch: int[source]#

0

gradient_check_parameters(transformer, latents, encoder_hidden_states, encoder_attention_mask, timesteps, target, eps=0.05, max_params_to_check=2000) → float[source]#

Verify gradients using finite differences for FSDP models with GRADIENT_CHECK_DTYPE. Uses standard tolerances for GRADIENT_CHECK_DTYPE precision.

initialize_training_pipeline(training_args: fastvideo.v1.fastvideo_args.TrainingArgs)[source]#
abstract initialize_validation_pipeline(training_args: fastvideo.v1.fastvideo_args.TrainingArgs)[source]#
setup_gradient_check(args, loader_iter, noise_scheduler, noise_random_generator) → float | None[source]#

Setup and perform gradient check on a fresh batch.

Parameters:

  • args – Training arguments

  • loader_iter – Data loader iterator

  • noise_scheduler – Noise scheduler for diffusion

  • noise_random_generator – Random number generator for noise

Returns:

Maximum gradient error or None if check is disabled/fails

Return type:

float or None

train_dataloader: torchdata.stateful_dataloader.StatefulDataLoader[source]#

None

train_loader_iter: Iterator[tuple[torch.Tensor, torch.Tensor, torch.Tensor, Dict[str, Any]]][source]#

None

abstract train_one_step(transformer, model_type, optimizer, lr_scheduler, loader, noise_scheduler, noise_random_generator, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm, weighting_scheme, logit_mean, logit_std, mode_scale)[source]#

Train one step of the model.

validation_pipeline: fastvideo.v1.pipelines.ComposedPipelineBase[source]#

None

fastvideo.v1.training.training_pipeline.logger[source]#

‘init_logger(…)’

Contents