--- title: cosmos2_5_training_pipeline ---

Cosmos 2.5 training pipeline (text-to-world, full fine-tuning + LoRA).

Follows the same structure as wan_training_pipeline.py. Key Cosmos 2.5 specifics: - Text embeddings are (B, seq, 100352) from Reason1 full-concat (28 layers × 3584 dim) - forward() takes padding_mask (B, 1, H, W) and optional fps int - Normalisation is handled by Cosmos25WanVAEWrapper (handles_latent_norm=True), so _normalize_dit_input is a no-op and stored latents are already normalised. - Timesteps are sigma values in [0, 1] (flow matching); (B,) is auto-expanded inside the model to (B, 1).

fastvideo.training.cosmos2_5_training_pipeline.get_local_torch_device

get_local_torch_device() -> torch.device

Return the torch device for the current rank.

fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs

FastVideoArgs(model_path: str, mode: ExecutionMode = ExecutionMode.INFERENCE, workload_type: WorkloadType = WorkloadType.T2V, distributed_executor_backend: str = 'mp', enable_streaming_ipc_queues: bool = False, ray_placement_group: PlacementGroup | None = None, ray_runtime_env: RuntimeEnv | None = None, inference_mode: bool = True, trust_remote_code: bool = False, revision: str | None = None, num_gpus: int = 1, tp_size: int = -1, sp_size: int = -1, hsdp_replicate_dim: int = 1, hsdp_shard_dim: int = -1, dist_timeout: int | None = None, pipeline_config: PipelineConfig = PipelineConfig(), preprocess_config: PreprocessConfig | None = None, lora_path: str | None = None, lora_nickname: str = 'default', lora_strength: float = 1.0, lora_target_modules: list[str] | None = None, output_type: str = 'pil', attention_backend: str | None = None, dit_cpu_offload: bool = True, use_fsdp_inference: bool = False, dit_layerwise_offload: bool = True, text_encoder_cpu_offload: bool = True, image_encoder_cpu_offload: bool = True, vae_cpu_offload: bool = True, pin_cpu_memory: bool = True, h3_sequential_load: bool | None = None, h3_encoder_split: bool | None = None, h3_encoder_nodes: int = -1, h3_encoder_workers: int = 0, video_decode_backend: str = 'h3-vae', taeh3_checkpoint: str | None = None, taeh3_chunk_size: int = 5, h3_vae_nvfp4_decoder: str | None = None, lazy_module_load: bool | None = None, vae_parallel_decode: bool = False, vae_parallel_encode: bool = False, vae_parallel_decode_strategy: str | None = None, enable_torch_compile: bool = False, regional_compile: bool = False, enable_torch_compile_text_encoder: bool = False, enable_torch_compile_vae: bool = False, enable_torch_compile_audio_vae: bool = False, torch_compile_kwargs: dict[str, Any] = dict(), torch_compile_kwargs_dit: dict[str, Any] = dict(), torch_compile_kwargs_text_encoder: dict[str, Any] = dict(), torch_compile_kwargs_vae: dict[str, Any] = dict(), torch_compile_kwargs_audio_vae: dict[str, Any] = dict(), inference_torch_compile: bool = False, h3_vsa_fp4: dict[str, Any] = dict(), disable_autocast: bool = False, VSA_sparsity: float | None = None, VSA_tile_size: int | None = None, moba_config_path: str | None = None, moba_config: dict[str, Any] = dict(), master_port: int | None = None, enable_stage_verification: bool = True, prompt_txt: str | None = None, ltx2_vae_tiling: bool | None = None, ltx2_vae_spatial_tile_size_in_pixels: int | None = None, ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None, ltx2_vae_temporal_tile_size_in_frames: int | None = None, ltx2_vae_temporal_tile_overlap_in_frames: int | None = None, ltx2_initial_latent_path: str | None = None, ltx2_audio_latent_path: str | None = None, refine_enabled: bool | None = None, refine_upsampler_path: str | None = None, refine_transformer_path: str | None = None, refine_lora_path: str | None = None, refine_num_inference_steps: int | None = None, refine_guidance_scale: float | None = None, refine_add_noise: bool | None = None, refine_noise_path: str | None = None, refine_audio_noise_path: str | None = None, ltx2_refine_enabled: bool = False, ltx2_refine_upsampler_path: str | None = None, ltx2_refine_transformer_path: str | None = None, ltx2_refine_lora_path: str | None = None, ltx2_refine_num_inference_steps: int = 3, ltx2_refine_guidance_scale: float = 1.0, ltx2_refine_add_noise: bool = True, ltx2_refine_noise_path: str | None = None, ltx2_refine_audio_noise_path: str | None = None, ltx2_legacy_native_noise_order: bool = False, ltx2_use_distilled_sigmas: bool = True, model_paths: dict[str, str] = dict(), model_loaded: dict[str, bool] = (lambda: {'transformer': True, 'vae': True, 'upsampler': True})(), override_text_encoder_safetensors: str | None = None, override_text_encoder_quant: QuantizationMethods = None, transformer_quant: Any | None = None, override_transformer_cls_name: str | None = None, init_weights_from_safetensors: str = '', init_weights_from_safetensors_2: str = '', override_pipeline_cls_name: str | None = None, boundary_ratio: float | None = 0.875)

fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.is_output_rank

is_output_rank: bool

Whether this process may materialize user-facing outputs.

This is runtime state assigned by an SPMD executor, not a public configuration field, so dataclass serialization and typed config parsing deliberately ignore it.

fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.add_cli_args

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.from_kwargs

from_kwargs(**kwargs: Any) -> FastVideoArgs

fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.check_fastvideo_args

check_fastvideo_args() -> None

Validate inference arguments for consistency

fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.finalize_device_offload_policy

finalize_device_offload_policy(device_id: int = 0) -> bool

Apply device-local memory policy, then resolve incompatible modes.

fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.disable_offload_on_unified_memory

disable_offload_on_unified_memory(device_id: int = 0, *, offload_flag: str | None = None) -> bool

Disable host offload after a worker has selected its device.

CUDA's unified-memory probe reads runtime device properties and may initialize a CUDA context. Callers must therefore use this only inside a device-owning process, after selecting and binding device_id. Returning the classification lets direct component-loader callers apply the same policy to explicit per-call overrides. When offload_flag is given, the return value says whether this policy covers that component role.

fastvideo.training.cosmos2_5_training_pipeline.TrainingArgs

TrainingArgs(model_path: str, mode: ExecutionMode = ExecutionMode.INFERENCE, workload_type: WorkloadType = WorkloadType.T2V, distributed_executor_backend: str = 'mp', enable_streaming_ipc_queues: bool = False, ray_placement_group: PlacementGroup | None = None, ray_runtime_env: RuntimeEnv | None = None, inference_mode: bool = True, trust_remote_code: bool = False, revision: str | None = None, num_gpus: int = 1, tp_size: int = -1, sp_size: int = -1, hsdp_replicate_dim: int = 1, hsdp_shard_dim: int = -1, dist_timeout: int | None = None, pipeline_config: PipelineConfig = PipelineConfig(), preprocess_config: PreprocessConfig | None = None, lora_path: str | None = None, lora_nickname: str = 'default', lora_strength: float = 1.0, lora_target_modules: list[str] | None = None, output_type: str = 'pil', attention_backend: str | None = None, dit_cpu_offload: bool = True, use_fsdp_inference: bool = False, dit_layerwise_offload: bool = True, text_encoder_cpu_offload: bool = True, image_encoder_cpu_offload: bool = True, vae_cpu_offload: bool = True, pin_cpu_memory: bool = True, h3_sequential_load: bool | None = None, h3_encoder_split: bool | None = None, h3_encoder_nodes: int = -1, h3_encoder_workers: int = 0, video_decode_backend: str = 'h3-vae', taeh3_checkpoint: str | None = None, taeh3_chunk_size: int = 5, h3_vae_nvfp4_decoder: str | None = None, lazy_module_load: bool | None = None, vae_parallel_decode: bool = False, vae_parallel_encode: bool = False, vae_parallel_decode_strategy: str | None = None, enable_torch_compile: bool = False, regional_compile: bool = False, enable_torch_compile_text_encoder: bool = False, enable_torch_compile_vae: bool = False, enable_torch_compile_audio_vae: bool = False, torch_compile_kwargs: dict[str, Any] = dict(), torch_compile_kwargs_dit: dict[str, Any] = dict(), torch_compile_kwargs_text_encoder: dict[str, Any] = dict(), torch_compile_kwargs_vae: dict[str, Any] = dict(), torch_compile_kwargs_audio_vae: dict[str, Any] = dict(), inference_torch_compile: bool = False, h3_vsa_fp4: dict[str, Any] = dict(), disable_autocast: bool = False, VSA_sparsity: float | None = None, VSA_tile_size: int | None = None, moba_config_path: str | None = None, moba_config: dict[str, Any] = dict(), master_port: int | None = None, enable_stage_verification: bool = True, prompt_txt: str | None = None, ltx2_vae_tiling: bool | None = None, ltx2_vae_spatial_tile_size_in_pixels: int | None = None, ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None, ltx2_vae_temporal_tile_size_in_frames: int | None = None, ltx2_vae_temporal_tile_overlap_in_frames: int | None = None, ltx2_initial_latent_path: str | None = None, ltx2_audio_latent_path: str | None = None, refine_enabled: bool | None = None, refine_upsampler_path: str | None = None, refine_transformer_path: str | None = None, refine_lora_path: str | None = None, refine_num_inference_steps: int | None = None, refine_guidance_scale: float | None = None, refine_add_noise: bool | None = None, refine_noise_path: str | None = None, refine_audio_noise_path: str | None = None, ltx2_refine_enabled: bool = False, ltx2_refine_upsampler_path: str | None = None, ltx2_refine_transformer_path: str | None = None, ltx2_refine_lora_path: str | None = None, ltx2_refine_num_inference_steps: int = 3, ltx2_refine_guidance_scale: float = 1.0, ltx2_refine_add_noise: bool = True, ltx2_refine_noise_path: str | None = None, ltx2_refine_audio_noise_path: str | None = None, ltx2_legacy_native_noise_order: bool = False, ltx2_use_distilled_sigmas: bool = True, model_paths: dict[str, str] = dict(), model_loaded: dict[str, bool] = (lambda: {'transformer': True, 'vae': True, 'upsampler': True})(), override_text_encoder_safetensors: str | None = None, override_text_encoder_quant: QuantizationMethods = None, transformer_quant: Any | None = None, override_transformer_cls_name: str | None = None, init_weights_from_safetensors: str = '', init_weights_from_safetensors_2: str = '', override_pipeline_cls_name: str | None = None, boundary_ratio: float | None = 0.875, pipeline_class: str = '', pipeline_module: str = '', data_path: str = '', dataloader_num_workers: int = 0, num_height: int = 0, num_width: int = 0, num_frames: int = 0, train_batch_size: int = 0, num_latent_t: int = 0, group_frame: bool = False, group_resolution: bool = False, pretrained_model_name_or_path: str = '', real_score_model_path: str = '', fake_score_model_path: str = '', ema_decay: float = 0.0, ema_start_step: int = 0, training_cfg_rate: float = 0.0, precondition_outputs: bool = False, validation_dataset_file: str = '', validation_preprocessed_path: str = '', validation_sampling_steps: str = '', validation_guidance_scale: str = '', validation_steps: float = 0.0, log_validation: bool = False, trackers: list[str] = list(), tracker_project_name: str = '', wandb_run_name: str = '', seed: int | None = None, output_dir: str = '', checkpoints_total_limit: int = 0, resume_from_checkpoint: str = '', num_train_epochs: int = 0, max_train_steps: int = 0, gradient_accumulation_steps: int = 0, learning_rate: float = 0.0, scale_lr: bool = False, lr_scheduler: str = 'constant', lr_warmup_steps: int = 0, max_grad_norm: float = 0.0, enable_gradient_checkpointing_type: str | None = None, selective_checkpointing: float = 0.0, mixed_precision: str = '', train_sp_batch_size: int = 0, fsdp_sharding_startegy: str = '', weighting_scheme: str = '', logit_mean: float = 0.0, logit_std: float = 1.0, mode_scale: float = 0.0, num_euler_timesteps: int = 0, lr_num_cycles: int = 0, lr_power: float = 0.0, min_lr_ratio: float = 0.5, not_apply_cfg_solver: bool = False, distill_cfg: float = 0.0, scheduler_type: str = '', linear_quadratic_threshold: float = 0.0, linear_range: float = 0.0, weight_decay: float = 0.0, betas: str = '0.9,0.999', use_ema: bool = False, multi_phased_distill_schedule: str = '', pred_decay_weight: float = 0.0, pred_decay_type: str = '', hunyuan_teacher_disable_cfg: bool = False, master_weight_type: str = '', VSA_decay_rate: float = 0.01, VSA_decay_interval_steps: int = 1, VSA_cache_tile_buf: bool = False, lora_rank: int | None = None, lora_alpha: int | None = None, lora_training: bool = False, ltx2_first_frame_conditioning_p: float = 0.1, generator_update_interval: int = 5, dfake_gen_update_ratio: int = 5, min_timestep_ratio: float = 0.2, max_timestep_ratio: float = 0.98, real_score_guidance_scale: float = 3.5, fake_score_learning_rate: float = 0.0, fake_score_lr_scheduler: str = 'constant', fake_score_betas: str = '0.9,0.999', training_state_checkpointing_steps: int = 0, weight_only_checkpointing_steps: int = 0, log_visualization: bool = False, visualization_steps: int = 0, simulate_generator_forward: bool = False, warp_denoising_step: bool = False, num_frame_per_block: int = 3, independent_first_frame: bool = False, enable_gradient_masking: bool = True, gradient_mask_last_n_frames: int = 21, same_step_across_blocks: bool = False, last_step_only: bool = False, context_noise: int = 0)

Training arguments. Inherits from FastVideoArgs and adds training-specific arguments. If there are any conflicts, the training arguments will take precedence.

fastvideo.training.cosmos2_5_training_pipeline.TrainingArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> TrainingArgs

fastvideo.training.cosmos2_5_training_pipeline.TrainingArgs.add_cli_args

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.training.cosmos2_5_training_pipeline.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler

FlowUniPCMultistepScheduler(num_train_timesteps: int = 1000, solver_order: int = 2, prediction_type: str = 'flow_prediction', shift: float | None = 1.0, use_dynamic_shifting = False, thresholding: bool = False, dynamic_thresholding_ratio: float = 0.995, sample_max_value: float = 1.0, predict_x0: bool = True, solver_type: str = 'bh2', lower_order_final: bool = True, disable_corrector: tuple = (), solver_p: SchedulerMixin = None, timestep_spacing: str = 'linspace', steps_offset: int = 0, final_sigmas_type: str | None = 'zero', **kwargs)

UniPCMultistepScheduler is a training-free framework designed for the fast sampling of diffusion models.

This model inherits from [SchedulerMixin] and [ConfigMixin]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving.

Parameters

num_train_timesteps: `int`, defaults to 1000

The number of diffusion steps to train the model.

solver_order: `int`, default `2`

The UniPC order which can be any positive integer. The effective order of accuracy is solver_order + 1 due to the UniC. It is recommended to use solver_order=2 for guided sampling, and solver_order=3 for unconditional sampling.

prediction_type: `str`, defaults to "flow_prediction"

Prediction type of the scheduler function; must be flow_prediction for this scheduler, which predicts the flow of the diffusion process.

thresholding: `bool`, defaults to `False`

Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such as Stable Diffusion.

dynamic_thresholding_ratio: `float`, defaults to 0.995

The ratio for the dynamic thresholding method. Valid only when thresholding=True.

sample_max_value: `float`, defaults to 1.0

The threshold value for dynamic thresholding. Valid only when thresholding=True and predict_x0=True.

predict_x0: `bool`, defaults to `True`

Whether to use the updating algorithm on the predicted x0.

solver_type: `str`, default `bh2`

Solver type for UniPC. It is recommended to use bh1 for unconditional sampling when steps < 10, and bh2 otherwise.

lower_order_final: `bool`, default `True`

Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.

disable_corrector: `list`, default `[]`

Decides which step to disable the corrector to mitigate the misalignment between epsilon_theta(x_t, c) and epsilon_theta(x_t^c, c) which can influence convergence for a large guidance scale. Corrector is usually disabled during the first few steps.

solver_p: `SchedulerMixin`, default `None`

Any other scheduler that if specified, the algorithm becomes solver_p + UniC.

use_karras_sigmas: `bool`, *optional*, defaults to `False`

Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If True, the sigmas are determined according to a sequence of noise levels {σi}.

use_exponential_sigmas: `bool`, *optional*, defaults to `False`

Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.

timestep_spacing: `str`, defaults to `"linspace"`

The way the timesteps should be scaled. Refer to Table 2 of the Common Diffusion Noise Schedules and Sample Steps are Flawed for more information.

steps_offset: `int`, defaults to 0

An offset added to the inference steps, as required by some model families.

final_sigmas_type: `str`, defaults to `"zero"`

The final sigma value for the noise schedule during the sampling process. If "sigma_min", the final sigma is the same as the last sigma in the training schedule. If zero, the final sigma is set to 0.

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.step_index

step_index

The index counter for current timestep. It will increase 1 after each scheduler step.

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.begin_index

begin_index

The index for the first timestep. It should be set from pipeline with set_begin_index method.

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.set_shift

set_shift(shift: float) -> None

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.set_begin_index

set_begin_index(begin_index: int = 0)

Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

Parameters

begin_index: `int`

The begin index for the scheduler.

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.set_timesteps

set_timesteps(num_inference_steps: int | None = None, device: str | torch.device = None, sigmas: list[float] | None = None, mu: float | None | None = None, shift: float | None | None = None, use_karras_sigmas: bool | None = None, use_kerras_sigma: bool | None = None)

Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (int): Total number of the spacing of the time steps. device (str or torch.device, optional): The device to which the timesteps should be moved to. If None, the timesteps are not moved.

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.time_shift

time_shift(mu: float, sigma: float, t: torch.Tensor)

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.convert_model_output

convert_model_output(model_output: torch.Tensor, *args, sample: torch.Tensor = None, **kwargs) -> torch.Tensor

Convert the model output to the corresponding type the UniPC algorithm needs.

Parameters

model_output: `torch.Tensor`

The direct output from the learned diffusion model.

timestep: `int`

The current discrete timestep in the diffusion chain.

sample: `torch.Tensor`

A current instance of a sample created by the diffusion process.

Returns

torch.Tensor

torch.Tensor: The converted model output.

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.multistep_uni_p_bh_update

multistep_uni_p_bh_update(model_output: torch.Tensor, *args, sample: torch.Tensor = None, order: int | None = None, **kwargs) -> torch.Tensor

One step for the UniP (B(h) version). Alternatively, self.solver_p is used if is specified.

Parameters

model_output: `torch.Tensor`

The direct output from the learned diffusion model at the current timestep.

prev_timestep: `int`

The previous discrete timestep in the diffusion chain.

sample: `torch.Tensor`

A current instance of a sample created by the diffusion process.

order: `int`

The order of UniP at this timestep (corresponds to the p in UniPC-p).

Returns

torch.Tensor

torch.Tensor: The sample tensor at the previous timestep.

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.multistep_uni_c_bh_update

multistep_uni_c_bh_update(this_model_output: torch.Tensor, *args, last_sample: torch.Tensor = None, this_sample: torch.Tensor = None, order: int | None = None, **kwargs) -> torch.Tensor

One step for the UniC (B(h) version).

Parameters

this_model_output: `torch.Tensor`

The model outputs at x_t.

this_timestep: `int`

The current timestep t.

last_sample: `torch.Tensor`

The generated sample before the last predictor x_{t-1}.

this_sample: `torch.Tensor`

The generated sample after the last predictor x_{t}.

order: `int`

The p of UniC-p at this step. The effective order of accuracy should be order + 1.

Returns

torch.Tensor

torch.Tensor: The corrected sample tensor at the current timestep.

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.index_for_timestep

index_for_timestep(timestep, schedule_timesteps = None) -> int

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.step

step(model_output: torch.Tensor, timestep: int | torch.Tensor, sample: torch.Tensor, return_dict: bool = True, generator = None) -> SchedulerOutput | tuple

Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep UniPC.

Parameters

model_output: `torch.Tensor`

The direct output from learned diffusion model.

timestep: `int`

The current discrete timestep in the diffusion chain.

sample: `torch.Tensor`

A current instance of a sample created by the diffusion process.

return_dict: `bool`

Whether or not to return a [~schedulers.scheduling_utils.SchedulerOutput] or tuple.

Returns

SchedulerOutput | tuple

[~schedulers.scheduling_utils.SchedulerOutput] or tuple: If return_dict is True, [~schedulers.scheduling_utils.SchedulerOutput] is returned, otherwise a tuple is returned where the first element is the sample tensor.

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.scale_model_input

scale_model_input(sample: torch.Tensor, *args, **kwargs) -> torch.Tensor

Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep.

Parameters

sample: `torch.Tensor`

The input sample.

Returns

torch.Tensor

torch.Tensor: A scaled input sample.

fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.add_noise

add_noise(original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor

fastvideo.training.cosmos2_5_training_pipeline.Cosmos2_5Pipeline

Cosmos2_5Pipeline(model_path: str, fastvideo_args: FastVideoArgs | TrainingArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None)

Cosmos 2.5 video generation pipeline.

fastvideo.training.cosmos2_5_training_pipeline.Cosmos2_5Pipeline.create_pipeline_stages

create_pipeline_stages(fastvideo_args: FastVideoArgs)

fastvideo.training.cosmos2_5_training_pipeline.TrainingBatch

TrainingBatch(current_timestep: int = 0, current_vsa_sparsity: float = 0.0, latents: torch.Tensor | None = None, raw_latent_shape: tuple[int, ...] | None = None, noise_latents: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, audio_latents: torch.Tensor | None = None, audio_noisy_model_input: torch.Tensor | None = None, audio_timesteps: torch.Tensor | None = None, audio_sigmas: torch.Tensor | None = None, audio_noise: torch.Tensor | None = None, audio_encoder_hidden_states: torch.Tensor | None = None, audio_encoder_attention_mask: torch.Tensor | None = None, conditioning_mask: torch.Tensor | None = None, preprocessed_image: torch.Tensor | None = None, image_embeds: torch.Tensor | None = None, image_latents: torch.Tensor | None = None, infos: list[dict[str, Any]] | None = None, mask_lat_size: torch.Tensor | None = None, trajectory_latents: torch.Tensor | None = None, trajectory_timesteps: torch.Tensor | None = None, noisy_model_input: torch.Tensor | None = None, timesteps: torch.Tensor | None = None, sigmas: torch.Tensor | None = None, noise: torch.Tensor | None = None, training_target: torch.Tensor | None = None, minimax_h3_layout: Any | None = None, attn_metadata_vsa: AttentionMetadata | None = None, attn_metadata: AttentionMetadata | None = None, input_kwargs: dict[str, Any] | None = None, loss: torch.Tensor | None = None, total_loss: float | None = None, grad_norm: float | None = None, encoder_hidden_states_neg: torch.Tensor | None = None, encoder_attention_mask_neg: torch.Tensor | None = None, conditional_dict: dict[str, Any] | None = None, unconditional_dict: dict[str, Any] | None = None, generator_loss: float = 0.0, fake_score_loss: float = 0.0, dmd_latent_vis_dict: dict[str, Any] = dict(), latent_vis_dict: dict[str, Any] = dict(), fake_score_latent_vis_dict: dict[str, Any] = dict(), minimax_h3_condition_video_rows: torch.Tensor | None = None, minimax_h3_condition_audio_rows: torch.Tensor | None = None, minimax_h3_dmd_layout: Any | None = None, minimax_h3_document_video_latents: tuple[torch.Tensor, ...] | None = None, minimax_h3_document_audio_latents: tuple[torch.Tensor, ...] | None = None, minimax_h3_text_attn_metadata: AttentionMetadata | None = None, sequence_packing_num_documents: int = 1, sequence_packing_num_tokens: int = 0, sequence_padded_num_tokens: int = 0, sequence_padding_ratio: float | None = None, sequence_packing_loss_scale: float = 1.0, sequence_packing_data_forced_ratio: float | None = None)

fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline

TrainingPipeline(model_path: str, fastvideo_args: TrainingArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None)

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.

fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.create_pipeline_stages

create_pipeline_stages(fastvideo_args: FastVideoArgs)

fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.set_schemas

set_schemas() -> None

fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.initialize_training_pipeline

initialize_training_pipeline(training_args: TrainingArgs)

fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.initialize_validation_pipeline

initialize_validation_pipeline(training_args: TrainingArgs)

fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.train_one_step

train_one_step(training_batch: TrainingBatch) -> TrainingBatch

fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.train

train() -> None

fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.visualize_intermediate_latents

visualize_intermediate_latents(training_batch: TrainingBatch, training_args: TrainingArgs, step: int)

Add visualization data to tracker logging and save frames to disk.

fastvideo.training.cosmos2_5_training_pipeline.Cosmos25TrainingPipeline

Cosmos25TrainingPipeline(model_path: str, fastvideo_args: TrainingArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None)

Training pipeline for Cosmos 2.5 (text-to-world).

Supports: - Full fine-tuning (all transformer parameters) - LoRA via the inherited LoRAPipeline mechanism (lora_param_names_mapping is set on Cosmos25Transformer3DModel)

Source
class Cosmos25TrainingPipeline(TrainingPipeline):
    """Training pipeline for Cosmos 2.5 (text-to-world).

    Supports:
    - Full fine-tuning (all transformer parameters)
    - LoRA via the inherited LoRAPipeline mechanism
      (lora_param_names_mapping is set on Cosmos25Transformer3DModel)
    """

    _required_config_modules = ["scheduler", "transformer", "vae"]

    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
        """Create the flow-matching scheduler with Cosmos 2.5's shift=5.0."""
        self.modules["scheduler"] = FlowUniPCMultistepScheduler(shift=fastvideo_args.pipeline_config.flow_shift)

    def initialize_validation_pipeline(self, training_args: TrainingArgs):
        """Build a full Cosmos2_5Pipeline that reuses the training transformer."""
        logger.info("Initializing Cosmos 2.5 validation pipeline...")
        args_copy = deepcopy(training_args)
        args_copy.inference_mode = True

        validation_pipeline = Cosmos2_5Pipeline.from_pretrained(
            training_args.model_path,
            args=args_copy,
            inference_mode=True,
            loaded_modules={
                "transformer": self.get_module("transformer"),
            },
            tp_size=training_args.tp_size,
            sp_size=training_args.sp_size,
            num_gpus=training_args.num_gpus,
            pin_cpu_memory=training_args.pin_cpu_memory,
            dit_cpu_offload=True,
        )
        self.validation_pipeline = validation_pipeline

    # ------------------------------------------------------------------
    # Cosmos 2.5-specific overrides
    # ------------------------------------------------------------------

    def _normalize_dit_input(self, training_batch: TrainingBatch) -> TrainingBatch:
        """Skip: Cosmos25WanVAEWrapper normalises latents during encoding.

        Pre-computed latents stored in the dataset are already normalised
        (per-channel mean/std applied inside the VAE wrapper), so applying a
        second normalisation here would corrupt the inputs.
        """
        return training_batch

    def _build_input_kwargs(self, training_batch: TrainingBatch) -> TrainingBatch:
        """Build the keyword arguments for Cosmos25Transformer3DModel.forward().

        Cosmos 2.5 forward() signature (T2W training subset):
            hidden_states : (B, 16, T, H, W)
            timestep      : (B,)  sigma values in [0, 1] — model expands to (B, 1)
            encoder_hidden_states : (B, seq, 100352)
            padding_mask  : (B, 1, H, W)  — all-ones for T2W (no spatial padding)
            fps           : int            — used only when rope_enable_fps_modulation=True
        """
        assert training_batch.noisy_model_input is not None
        assert training_batch.encoder_hidden_states is not None
        assert training_batch.sigmas is not None

        noisy = training_batch.noisy_model_input  # (B, 16, T, H, W)
        batch_size, _c, _t, height, width = noisy.shape
        device = get_local_torch_device()

        # sigmas stored as (B, 1, 1, 1, 1) after get_sigmas; flatten to (B,)
        sigma_vals = training_batch.sigmas.view(batch_size).to(dtype=torch.bfloat16, device=device)

        # T2W: no spatial padding, all-ones padding mask (16 latent + 1 padding).
        padding_mask = torch.ones(
            batch_size,
            1,
            height,
            width,
            device=device,
            dtype=torch.bfloat16,
        )

        # T2W: no inpainting, all-zeros condition mask (adds 1 channel → total 18).
        condition_mask = torch.zeros(
            batch_size,
            1,
            _t,
            height,
            width,
            device=device,
            dtype=torch.bfloat16,
        )

        training_batch.input_kwargs = {
            "hidden_states": noisy,
            "encoder_hidden_states": training_batch.encoder_hidden_states,
            # (B,) sigma in [0, 1]; model auto-expands to (B, 1) at line 901
            "timestep": sigma_vals,
            "condition_mask": condition_mask,
            "padding_mask": padding_mask,
            # rope_enable_fps_modulation=False for the 2B checkpoint, so fps
            # has no effect on the embeddings, but the argument is still forwarded.
            "fps": 24,
        }
        return training_batch

fastvideo.training.cosmos2_5_training_pipeline.Cosmos25TrainingPipeline.initialize_pipeline

initialize_pipeline(fastvideo_args: FastVideoArgs)

Create the flow-matching scheduler with Cosmos 2.5's shift=5.0.

Source
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
    """Create the flow-matching scheduler with Cosmos 2.5's shift=5.0."""
    self.modules["scheduler"] = FlowUniPCMultistepScheduler(shift=fastvideo_args.pipeline_config.flow_shift)

fastvideo.training.cosmos2_5_training_pipeline.Cosmos25TrainingPipeline.initialize_validation_pipeline

initialize_validation_pipeline(training_args: TrainingArgs)

Build a full Cosmos2_5Pipeline that reuses the training transformer.

Source
def initialize_validation_pipeline(self, training_args: TrainingArgs):
    """Build a full Cosmos2_5Pipeline that reuses the training transformer."""
    logger.info("Initializing Cosmos 2.5 validation pipeline...")
    args_copy = deepcopy(training_args)
    args_copy.inference_mode = True

    validation_pipeline = Cosmos2_5Pipeline.from_pretrained(
        training_args.model_path,
        args=args_copy,
        inference_mode=True,
        loaded_modules={
            "transformer": self.get_module("transformer"),
        },
        tp_size=training_args.tp_size,
        sp_size=training_args.sp_size,
        num_gpus=training_args.num_gpus,
        pin_cpu_memory=training_args.pin_cpu_memory,
        dit_cpu_offload=True,
    )
    self.validation_pipeline = validation_pipeline

fastvideo.training.cosmos2_5_training_pipeline.main

main(args) -> None

fastvideo.training.cosmos2_5_training_pipeline.build_parser

build_parser(pipeline_class: str | None = None, pipeline_module: str | None = None) -> FlexibleArgumentParser

Build the training-runner CLI parser.

When pipeline_class/pipeline_module are given they become optional defaults, so the deprecated per-pipeline __main__ blocks keep working without the flags.