Skip to content

loader

Modules

fastvideo.models.loader.benchmarks

Modules

fastvideo.models.loader.benchmarks.benchmark_weight_loading

Benchmark for model weight loading speed.

Measures the time to load model weights from safetensors files using different strategies (CPU vs GPU, broadcast vs independent).

Usage (single GPU): python fastvideo/models/loader/benchmarks/benchmark_weight_loading.py --model-path /path/to/model

Usage (multi-GPU, e.g. 4 GPUs): torchrun --nproc_per_node=4 fastvideo/models/loader/benchmarks/benchmark_weight_loading.py --model-path /path/to/model

Functions:
fastvideo.models.loader.benchmarks.benchmark_weight_loading.benchmark_loading 
benchmark_loading(files: list[str], to_cpu: bool, broadcast: bool, warmup: int, repeats: int, label: str) -> None

Run the weight loading benchmark and print results.

Source code in fastvideo/models/loader/benchmarks/benchmark_weight_loading.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def benchmark_loading(
    files: list[str],
    to_cpu: bool,
    broadcast: bool,
    warmup: int,
    repeats: int,
    label: str,
) -> None:
    """Run the weight loading benchmark and print results."""
    rank = dist.get_rank() if dist.is_initialized() else 0
    node_group = get_node_group()

    # Count total params and bytes on first pass
    total_params = 0
    total_bytes = 0
    for name, tensor in safetensors_weights_iterator(
            files, to_cpu=to_cpu, broadcast=broadcast):
        total_params += 1
        total_bytes += tensor.nelement() * tensor.element_size()

    if rank == 0:
        logger.info("[%s] %d tensors, %.2f GB total",
                    label, total_params, total_bytes / 1e9)

    # Warmup
    for _ in range(warmup):
        for _ in safetensors_weights_iterator(
                files, to_cpu=to_cpu, broadcast=broadcast):
            pass
        if dist.is_initialized():
            dist.barrier()

    # Timed runs
    times = []
    for i in range(repeats):
        if dist.is_initialized():
            dist.barrier()
        torch.cuda.synchronize() if torch.cuda.is_available() else None

        t0 = time.perf_counter()
        for _ in safetensors_weights_iterator(
                files, to_cpu=to_cpu, broadcast=broadcast):
            pass
        torch.cuda.synchronize() if torch.cuda.is_available() else None

        if dist.is_initialized():
            dist.barrier()
        elapsed = time.perf_counter() - t0
        times.append(elapsed)

    if rank == 0:
        avg = sum(times) / len(times)
        best = min(times)
        throughput = total_bytes / avg / 1e9
        logger.info(
            "[%s] avg %.3fs | best %.3fs | throughput %.2f GB/s "
            "(over %d runs, %d warmup, %d GPU(s), node_size=%d)",
            label, avg, best, throughput,
            repeats, warmup, node_group.world_size, node_group.world_size,
        )
fastvideo.models.loader.benchmarks.benchmark_weight_loading_comparison

A/B benchmark: independent-read vs rank-0-broadcast weight loading.

Compares two strategies
  • "before" (independent): every rank reads safetensors from disk to GPU
  • "after" (broadcast): rank 0 reads from disk, broadcasts to other ranks
Usage
1 GPU

torchrun --nproc_per_node=1 fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py --model-path /path/to/model --subfolder transformer

2 GPUs

torchrun --nproc_per_node=2 fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py --model-path /path/to/model --subfolder transformer

4 GPUs

torchrun --nproc_per_node=4 fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py --model-path /path/to/model --subfolder transformer

Functions:
fastvideo.models.loader.benchmarks.benchmark_weight_loading_comparison.load_broadcast 
load_broadcast(files: list[str], device: str, node_group, async_op: bool = False)

After-PR behavior: rank 0 reads from disk, broadcasts to other ranks.

Source code in fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def load_broadcast(files: list[str], device: str, node_group,
                   async_op: bool = False):
    """After-PR behavior: rank 0 reads from disk, broadcasts to other ranks."""
    local_rank = node_group.local_rank
    handles = []
    for st_file in files:
        with safe_open(st_file, framework="pt", device=device) as f:
            for name in f:
                if local_rank == 0:
                    param = f.get_tensor(name)
                else:
                    sl = f.get_slice(name)
                    shape = sl.get_shape()
                    dtype = SAFETENSORS_TO_TORCH_DTYPE[sl.get_dtype()]
                    param = torch.empty(shape, device=device, dtype=dtype)
                if node_group.world_size > 1:
                    group = node_group.device_group
                    if async_op:
                        handle = dist.broadcast(
                            param,
                            src=dist.get_global_rank(group, 0),
                            async_op=True,
                            group=group,
                        )
                        handles.append(handle)
                    else:
                        dist.broadcast(
                            param,
                            src=dist.get_global_rank(group, 0),
                            group=group,
                        )
                yield name, param
        if async_op:
            for handle in handles:
                handle.wait()
            handles.clear()
fastvideo.models.loader.benchmarks.benchmark_weight_loading_comparison.load_independent 
load_independent(files: list[str], device: str)

Before-PR behavior: every rank reads every tensor from disk to GPU.

Source code in fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py 
41
42
43
44
45
46
47
def load_independent(files: list[str], device: str):
    """Before-PR behavior: every rank reads every tensor from disk to GPU."""
    for st_file in files:
        with safe_open(st_file, framework="pt", device=device) as f:
            for name in f:
                param = f.get_tensor(name)
                yield name, param

fastvideo.models.loader.component_loader

Classes

fastvideo.models.loader.component_loader.AudioDecoderLoader 
AudioDecoderLoader(device=None)

Bases: ComponentLoader

Loader for LTX-2 audio decoder (audio_vae component).

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device
fastvideo.models.loader.component_loader.ComponentLoader 
ComponentLoader(device=None)

Bases: ABC

Base class for loading a specific type of model component.

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device
Methods:
fastvideo.models.loader.component_loader.ComponentLoader.for_module_type classmethod 
for_module_type(module_type: str, transformers_or_diffusers: str) -> ComponentLoader

Factory method to create a component loader for a specific module type.

Parameters:

Name Type Description Default
module_type str

Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler")

 required
transformers_or_diffusers str

Whether the module is from transformers or diffusers

 required

Returns:

Type Description
ComponentLoader

A component loader for the specified module type
Source code in fastvideo/models/loader/component_loader.py 
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
@classmethod
def for_module_type(
    cls, module_type: str, transformers_or_diffusers: str
) -> "ComponentLoader":
    """
    Factory method to create a component loader for a specific module type.

    Args:
        module_type: Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler")
        transformers_or_diffusers: Whether the module is from transformers or diffusers

    Returns:
        A component loader for the specified module type
    """
    # Map of module types to their loader classes and expected library
    module_loaders = {
        "scheduler": (SchedulerLoader, "diffusers"),
        "transformer": (TransformerLoader, "diffusers"),
        "sr_transformer": (TransformerLoader, "diffusers"),
        "transformer_2": (TransformerLoader, "diffusers"),
        "transformer_3": (TransformerLoader, "diffusers"),
        "vae": (VAELoader, "diffusers"),
        "light_vae": (VAELoader, "diffusers"),
        "audio_vae": (AudioDecoderLoader, "diffusers"),
        "audio_decoder": (AudioDecoderLoader, "diffusers"),
        "vocoder": (VocoderLoader, "diffusers"),
        "text_encoder": (TextEncoderLoader, "transformers"),
        "text_encoder_2": (TextEncoderLoader, "transformers"),
        "text_encoder_3": (TextEncoderLoader, "transformers"),
        "tokenizer": (TokenizerLoader, "transformers"),
        "tokenizer_2": (TokenizerLoader, "transformers"),
        "tokenizer_3": (TokenizerLoader, "transformers"),
        "image_processor": (ImageProcessorLoader, "transformers"),
        "feature_extractor": (ImageProcessorLoader, "transformers"),
        "image_encoder": (ImageEncoderLoader, "transformers"),
        "upsampler": (UpsamplerLoader, "diffusers"),
        "upsampler_2": (UpsamplerLoader, "diffusers"),
        # Stable Audio's `StableAudioMultiConditioner` bundles T5 +
        # NumberConditioners; not a pure text encoder, so it gets
        # its own loader.
        "conditioner": (ConditionerLoader, "fastvideo"),
        # LTX-2 spatial / temporal upsamplers — share the
        # UpsamplerLoader path with the upsampler/upsampler_2 keys
        # so the SR pipeline picks up real weights instead of the
        # generic config-only loader.
        "spatial_upsampler": (UpsamplerLoader, "diffusers"),
        "temporal_upsampler": (UpsamplerLoader, "diffusers"),
    }

    if module_type in module_loaders:
        loader_cls, expected_library = module_loaders[module_type]
        # Allow fastvideo.* libraries for custom implementations (e.g. Cosmos2_5Pipeline)
        # that aren't available in diffusers/transformers yet
        is_fastvideo_module = transformers_or_diffusers.startswith("fastvideo.")
        if not is_fastvideo_module:
            # Assert that the library matches what's expected for this module type
            assert transformers_or_diffusers == expected_library, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}"
        return loader_cls()

    # For unknown module types, use a generic loader
    logger.warning(
        "No specific loader found for module type: %s. Using generic loader.",
        module_type,
    )
    return GenericComponentLoader(transformers_or_diffusers)
fastvideo.models.loader.component_loader.ComponentLoader.load abstractmethod 
load(model_path: str, fastvideo_args: FastVideoArgs)

Load the component based on the model path, architecture, and inference args.

Parameters:

Name Type Description Default
model_path str

Path to the component model

 required
fastvideo_args FastVideoArgs

FastVideoArgs

 required

Returns:

Type Description

The loaded component
Source code in fastvideo/models/loader/component_loader.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
@abstractmethod
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """
    Load the component based on the model path, architecture, and inference args.

    Args:
        model_path: Path to the component model
        fastvideo_args: FastVideoArgs

    Returns:
        The loaded component
    """
    raise NotImplementedError
fastvideo.models.loader.component_loader.ConditionerLoader 
ConditionerLoader(device=None)

Bases: ComponentLoader

Loader for multi-conditioner components (e.g. Stable Audio's StableAudioMultiConditioner, which bundles T5 + NumberConditioners and is neither a pure text encoder nor a Diffusers-shaped module). Reads <subfolder>/config.json to resolve the class via ModelRegistry, instantiates with no args (the class pulls its own defaults from its FastVideo config), then loads diffusion_pytorch_model.safetensors non-strictly so externally fetched sub-encoders (T5) don't trip the missing-key check.

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device
fastvideo.models.loader.component_loader.GenericComponentLoader 
GenericComponentLoader(library='transformers')

Bases: ComponentLoader

Generic loader for components that don't have a specific loader.

Source code in fastvideo/models/loader/component_loader.py 
1207
1208
1209
def __init__(self, library="transformers") -> None:
    super().__init__()
    self.library = library
Methods:
fastvideo.models.loader.component_loader.GenericComponentLoader.load 
load(model_path: str, fastvideo_args: FastVideoArgs)

Load a generic component based on the model path, and inference args.

Source code in fastvideo/models/loader/component_loader.py 
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load a generic component based on the model path, and inference args."""
    logger.warning(
        "Using generic loader for %s with library %s",
        model_path,
        self.library,
    )

    if self.library == "transformers":
        from transformers import AutoModel

        model = AutoModel.from_pretrained(
            model_path,
            trust_remote_code=fastvideo_args.trust_remote_code,
            revision=fastvideo_args.revision,
        )
        logger.info(
            "Loaded generic transformers model: %s",
            model.__class__.__name__,
        )
        return model
    elif self.library == "diffusers":
        logger.warning(
            "Generic loading for diffusers components is not fully implemented"
        )

        model_config = get_diffusers_config(model=model_path)
        logger.info("Diffusers Model config: %s", model_config)
        # This is a placeholder - in a real implementation, you'd need to handle this properly
        return None
    else:
        raise ValueError(f"Unsupported library: {self.library}")
fastvideo.models.loader.component_loader.ImageEncoderLoader 
ImageEncoderLoader(device=None)

Bases: TextEncoderLoader

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device
Methods:
fastvideo.models.loader.component_loader.ImageEncoderLoader.load 
load(model_path: str, fastvideo_args: FastVideoArgs)

Load the text encoders based on the model path, and inference args.

Source code in fastvideo/models/loader/component_loader.py 
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the text encoders based on the model path, and inference args."""
    # model_config: PretrainedConfig = get_hf_config(
    #     model=model_path,
    #     trust_remote_code=fastvideo_args.trust_remote_code,
    #     revision=fastvideo_args.revision,
    #     model_override_args=None,
    # )
    with open(os.path.join(model_path, "config.json")) as f:
        model_config = json.load(f)
    model_config.pop("_name_or_path", None)
    model_config.pop("transformers_version", None)
    model_config.pop("torch_dtype", None)
    model_config.pop("model_type", None)
    logger.info("HF Model config: %s", model_config)

    encoder_config = fastvideo_args.pipeline_config.image_encoder_config
    encoder_config.update_model_arch(model_config)

    from fastvideo.platforms import current_platform

    if fastvideo_args.image_encoder_cpu_offload:
        target_device = (
            torch.device("mps")
            if current_platform.is_mps()
            else torch.device("cpu")
        )
    else:
        target_device = get_local_torch_device()
    # TODO(will): add support for other dtypes
    return self.load_model(
        model_path,
        encoder_config,
        target_device,
        fastvideo_args,
        fastvideo_args.pipeline_config.image_encoder_precision,
    )
fastvideo.models.loader.component_loader.ImageProcessorLoader 
ImageProcessorLoader(device=None)

Bases: ComponentLoader

Loader for image processor.

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device
Methods:
fastvideo.models.loader.component_loader.ImageProcessorLoader.load 
load(model_path: str, fastvideo_args: FastVideoArgs)

Load the image processor based on the model path, and inference args.

Source code in fastvideo/models/loader/component_loader.py 
514
515
516
517
518
519
520
521
522
523
524
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the image processor based on the model path, and inference args."""
    logger.info("Loading image processor from %s", model_path)

    image_processor = AutoImageProcessor.from_pretrained(
        model_path,
    )
    logger.info(
        "Loaded image processor: %s", image_processor.__class__.__name__
    )
    return image_processor
fastvideo.models.loader.component_loader.PipelineComponentLoader

Utility class for loading pipeline components. This replaces the chain of if-else statements in load_pipeline_module.

Methods:
fastvideo.models.loader.component_loader.PipelineComponentLoader.load_module staticmethod 
load_module(module_name: str, component_model_path: str, transformers_or_diffusers: str, fastvideo_args: FastVideoArgs)

Load a pipeline module.

Parameters:

Name Type Description Default
module_name str

Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler")

 required
component_model_path str

Path to the component model

 required
transformers_or_diffusers str

Whether the module is from transformers or diffusers

 required
pipeline_args

Inference arguments

 required

Returns:

Type Description

The loaded module
Source code in fastvideo/models/loader/component_loader.py 
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
@staticmethod
def load_module(
    module_name: str,
    component_model_path: str,
    transformers_or_diffusers: str,
    fastvideo_args: FastVideoArgs,
):
    """
    Load a pipeline module.

    Args:
        module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler")
        component_model_path: Path to the component model
        transformers_or_diffusers: Whether the module is from transformers or diffusers
        pipeline_args: Inference arguments

    Returns:
        The loaded module
    """
    logger.info(
        "Loading %s using %s from %s",
        module_name,
        transformers_or_diffusers,
        component_model_path,
    )

    # Get the appropriate loader for this module type
    loader = ComponentLoader.for_module_type(
        module_name, transformers_or_diffusers
    )

    # Load the module
    return loader.load(component_model_path, fastvideo_args)
fastvideo.models.loader.component_loader.SchedulerLoader 
SchedulerLoader(device=None)

Bases: ComponentLoader

Loader for scheduler.

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device
Methods:
fastvideo.models.loader.component_loader.SchedulerLoader.load 
load(model_path: str, fastvideo_args: FastVideoArgs)

Load the scheduler based on the model path, and inference args.

Source code in fastvideo/models/loader/component_loader.py 
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the scheduler based on the model path, and inference args."""
    config = get_diffusers_config(model=model_path)

    class_name = config.pop("_class_name")
    assert class_name is not None, (
        "Model config does not contain a _class_name attribute. Only diffusers format is supported."
    )

    scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)

    scheduler = scheduler_cls(**config)
    if fastvideo_args.pipeline_config.flow_shift is not None:
        scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift)
    return scheduler
fastvideo.models.loader.component_loader.TextEncoderLoader 
TextEncoderLoader(device=None)

Bases: ComponentLoader

Loader for text encoders.

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device
Classes
fastvideo.models.loader.component_loader.TextEncoderLoader.Source dataclass 
Source(model_or_path: str, prefix: str = '', fall_back_to_pt: bool = True, allow_patterns_overrides: list[str] | None = None)

A source for weights.

Attributes
fastvideo.models.loader.component_loader.TextEncoderLoader.Source.allow_patterns_overrides class-attribute instance-attribute 
allow_patterns_overrides: list[str] | None = None

If defined, weights will load exclusively using these patterns.

fastvideo.models.loader.component_loader.TextEncoderLoader.Source.fall_back_to_pt class-attribute instance-attribute 
fall_back_to_pt: bool = True

Whether .pt weights can be used.

fastvideo.models.loader.component_loader.TextEncoderLoader.Source.model_or_path instance-attribute 
model_or_path: str

The model ID or path.

fastvideo.models.loader.component_loader.TextEncoderLoader.Source.prefix class-attribute instance-attribute 
prefix: str = ''

A prefix to prepend to all weights.

Methods:
fastvideo.models.loader.component_loader.TextEncoderLoader.load 
load(model_path: str, fastvideo_args: FastVideoArgs)

Load the text encoders based on the model path, and inference args.

Source code in fastvideo/models/loader/component_loader.py 
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the text encoders based on the model path, and inference args."""
    # model_config: PretrainedConfig = get_hf_config(
    #     model=model_path,
    #     trust_remote_code=fastvideo_args.trust_remote_code,
    #     revision=fastvideo_args.revision,
    #     model_override_args=None,
    # )
    model_config = get_diffusers_config(model=model_path)
    model_config.pop("_name_or_path", None)
    model_config.pop("transformers_version", None)
    model_config.pop("model_type", None)
    model_config.pop("tokenizer_class", None)
    model_config.pop("torch_dtype", None)
    repo_root = os.path.dirname(model_path)
    index_path = os.path.join(repo_root, "model_index.json")
    gemma_path = ""
    gemma_path_from_candidate = False
    if os.path.isfile(index_path):
        try:
            with open(index_path, encoding="utf-8") as f:
                model_index = json.load(f)
            gemma_path = model_index.get("gemma_model_path", "")
        except json.JSONDecodeError:
            gemma_path = ""
    if not gemma_path:
        candidate = os.path.normpath(os.path.join(model_path, "gemma"))
        if os.path.isdir(candidate):
            gemma_path = candidate
            gemma_path_from_candidate = True
            model_config["gemma_model_path"] = gemma_path
    if gemma_path and not gemma_path_from_candidate:
        if not os.path.isabs(gemma_path):
            model_config["gemma_model_path"] = os.path.normpath(
                os.path.join(repo_root, gemma_path)
            )
    transformer_config_path = os.path.join(
        repo_root, "transformer", "config.json"
    )
    if os.path.isfile(transformer_config_path):
        try:
            with open(transformer_config_path, encoding="utf-8") as f:
                transformer_config = json.load(f)
            if (
                "connector_double_precision_rope" not in model_config
                or not model_config["connector_double_precision_rope"]
            ):
                if transformer_config.get("double_precision_rope") is True:
                    model_config["connector_double_precision_rope"] = True
            if "connector_rope_type" not in model_config:
                rope_type = transformer_config.get("rope_type")
                if rope_type is not None:
                    model_config["connector_rope_type"] = rope_type
        except json.JSONDecodeError:
            pass
    logger.info("HF Model config: %s", model_config)

    base = os.path.basename(os.path.normpath(model_path))
    idx = 0
    if base.startswith("text_encoder_"):
        try:
            idx = int(base.split("_")[-1]) - 1
        except Exception:
            idx = 0
    encoder_configs = fastvideo_args.pipeline_config.text_encoder_configs
    encoder_precisions = fastvideo_args.pipeline_config.text_encoder_precisions
    if idx < 0 or idx >= len(encoder_configs):
        raise IndexError(
            f"text encoder index {idx} out of range for text_encoder_configs (len={len(encoder_configs)}), model_path={model_path}"
        )
    encoder_config = encoder_configs[idx]
    encoder_config.update_model_arch(model_config)
    if idx < 0 or idx >= len(encoder_precisions):
        raise IndexError(
            f"text encoder index {idx} out of range for text_encoder_precisions (len={len(encoder_precisions)}), model_path={model_path}"
        )
    encoder_precision = encoder_precisions[idx]

    target_device = get_local_torch_device()
    # TODO(will): add support for other dtypes
    return self.load_model(
        model_path,
        encoder_config,
        target_device,
        fastvideo_args,
        encoder_precision,
        use_text_encoder_override=True,
    )
fastvideo.models.loader.component_loader.TokenizerLoader 
TokenizerLoader(device=None)

Bases: ComponentLoader

Loader for tokenizers.

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device
Methods:
fastvideo.models.loader.component_loader.TokenizerLoader.load 
load(model_path: str, fastvideo_args: FastVideoArgs)

Load the tokenizer based on the model path, and inference args.

Source code in fastvideo/models/loader/component_loader.py 
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the tokenizer based on the model path, and inference args."""
    logger.info("Loading tokenizer from %s", model_path)
    resolved_model_path = model_path

    # LTX2 checkpoints may not ship a top-level tokenizer/ directory.
    # In that case, tokenizer assets live under text_encoder/gemma/.
    if not os.path.isdir(resolved_model_path):
        ltx2_gemma_path = os.path.normpath(
            os.path.join(resolved_model_path, "..", "text_encoder",
                         "gemma"))
        if os.path.isdir(ltx2_gemma_path):
            logger.info(
                "Tokenizer directory %s missing; falling back to %s",
                resolved_model_path,
                ltx2_gemma_path,
            )
            resolved_model_path = ltx2_gemma_path

    # Cosmos2.5 stores an AutoProcessor config in `tokenizer/config.json` (not a tokenizer
    # config). Use its `_name_or_path` (e.g. Qwen/Qwen2.5-VL-7B-Instruct) as the source.
    tokenizer_cfg_path = os.path.join(resolved_model_path, "config.json")
    if os.path.exists(tokenizer_cfg_path):
        try:
            with open(tokenizer_cfg_path, "r") as f:
                tokenizer_cfg = json.load(f)
            if isinstance(tokenizer_cfg, dict) and (
                tokenizer_cfg.get("_class_name") == "AutoProcessor"
                or "processor_type" in tokenizer_cfg
            ):
                src = tokenizer_cfg.get("_name_or_path", "")
                if isinstance(src, str) and src.strip():
                    processor = AutoProcessor.from_pretrained(
                        src.strip(),
                        trust_remote_code=True,
                    )
                    logger.info(
                        "Loaded tokenizer/processor from %s: %s",
                        src,
                        processor.__class__.__name__,
                    )
                    return processor
        except Exception:
            # If parsing fails, fall through to AutoTokenizer below.
            pass

    # Only Flux2 full's Mistral3 (require_processor=True) must load via
    # AutoProcessor. Gate the processor_config.json shortcut on that flag so
    # existing encoders (e.g. HunyuanVideo 1.5 / Qwen2.5-VL) stay on the
    # historical AutoTokenizer path below even if their tokenizer dir happens
    # to ship a processor_config.json.
    require_processor = False
    if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"):
        try:
            require_processor = any(
                getattr(getattr(cfg, "arch_config", None), "require_processor", False)
                for cfg in fastvideo_args.pipeline_config.text_encoder_configs)
        except Exception:
            require_processor = False

    if require_processor and os.path.exists(os.path.join(resolved_model_path, "processor_config.json")):
        processor = AutoProcessor.from_pretrained(
            resolved_model_path,
            local_files_only=os.path.isdir(resolved_model_path),
            trust_remote_code=fastvideo_args.trust_remote_code,
        )
        logger.info(
            "Loaded tokenizer/processor from %s: %s",
            resolved_model_path,
            processor.__class__.__name__,
        )
        return processor

    try:
        tokenizer = AutoTokenizer.from_pretrained(
            resolved_model_path,  # "<path to model>/tokenizer"
            # in v0, this was same string as encoder_name "ClipTextModel"
            # TODO(will): pass these tokenizer kwargs from inference args? Maybe
            # other method of config?
            local_files_only=os.path.isdir(resolved_model_path),
        )
    except (OSError, ValueError):
        tokenizer = AutoProcessor.from_pretrained(
            resolved_model_path,
            local_files_only=os.path.isdir(resolved_model_path),
            trust_remote_code=fastvideo_args.trust_remote_code,
        )
        logger.info(
            "Loaded tokenizer/processor from %s: %s",
            resolved_model_path,
            tokenizer.__class__.__name__,
        )
        return tokenizer
    padding_side = None
    if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"):
        try:
            arch_config = fastvideo_args.pipeline_config.text_encoder_configs[
                0
            ].arch_config
            padding_side = getattr(arch_config, "padding_side", None)
        except Exception:
            padding_side = None
    if padding_side:
        tokenizer.padding_side = padding_side
    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
        tokenizer.pad_token = tokenizer.eos_token
    logger.info("Loaded tokenizer: %s", tokenizer.__class__.__name__)
    return tokenizer
fastvideo.models.loader.component_loader.TransformerLoader 
TransformerLoader(device=None)

Bases: ComponentLoader

Loader for transformer.

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device
Methods:
fastvideo.models.loader.component_loader.TransformerLoader.load 
load(model_path: str, fastvideo_args: FastVideoArgs)

Load the transformer based on the model path, and inference args.

Source code in fastvideo/models/loader/component_loader.py 
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the transformer based on the model path, and inference args."""
    config = get_diffusers_config(model=model_path)
    hf_config = deepcopy(config)
    cls_name = config.pop("_class_name")
    config.pop("_name_or_path", None)
    if cls_name is None:
        raise ValueError(
            "Model config does not contain a _class_name attribute. "
            "Only diffusers format is supported."
        )

    logger.info("transformer cls_name: %s", cls_name)
    if fastvideo_args.override_transformer_cls_name is not None:
        cls_name = fastvideo_args.override_transformer_cls_name
        logger.info("Overriding transformer cls_name to %s", cls_name)

    fastvideo_args.model_paths["transformer"] = model_path

    # Config from Diffusers supersedes fastvideo's model config
    dit_config = deepcopy(fastvideo_args.pipeline_config.dit_config)
    dit_config.update_model_arch(config)

    model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)

    # Find all safetensors files
    safetensors_list = glob.glob(
        os.path.join(str(model_path), "*.safetensors")
    )
    if not safetensors_list:
        raise ValueError(f"No safetensors files found in {model_path}")

    # arch_config can infer architecture from weight keys (e.g. Flux2 layer counts)
    update_fn = getattr(dit_config.arch_config, "update_from_weight_keys", None)
    if callable(update_fn):
        weight_keys = _collect_safetensors_keys(safetensors_list)
        update_fn(weight_keys)

    # Check if we should use custom initialization weights
    custom_weights_path = getattr(
        fastvideo_args, "init_weights_from_safetensors", None
    )
    use_custom_weights = (
        custom_weights_path
        and os.path.exists(custom_weights_path)
        and not hasattr(fastvideo_args, "_loading_teacher_critic_model")
    )

    if use_custom_weights:
        if "transformer_2" in model_path:
            custom_weights_path = getattr(
                fastvideo_args, "init_weights_from_safetensors_2", None
            )
        assert custom_weights_path is not None, (
            "Custom initialization weights must be provided"
        )
        if os.path.isdir(custom_weights_path):
            safetensors_list = glob.glob(
                os.path.join(str(custom_weights_path), "*.safetensors")
            )
        else:
            assert custom_weights_path.endswith(".safetensors"), (
                "Custom initialization weights must be a safetensors file"
            )
            safetensors_list = [custom_weights_path]

    logger.info(
        "Loading model from %s safetensors files: %s",
        len(safetensors_list),
        safetensors_list,
    )

    default_dtype = PRECISION_TO_TYPE[
        fastvideo_args.pipeline_config.dit_precision
    ]

    # Load the model using FSDP loader
    logger.info("Loading model from %s, default_dtype: %s", cls_name,
                default_dtype)
    assert fastvideo_args.hsdp_shard_dim is not None
    # Cosmos2.5 checkpoints can include extra entries not present in the
    # instantiated model (e.g. pos_embedder ranges / *_extra_state). Load
    # non-strictly for Cosmos2.5 only; keep upstream strict behavior for others.
    strict_load = not (
        cls_name.startswith("Cosmos25")
        or cls_name == "Cosmos25Transformer3DModel"
        or getattr(fastvideo_args.pipeline_config, "prefix", "") == "Cosmos25"
    )
    model = maybe_load_fsdp_model(
        model_cls=model_cls,
        init_params={"config": dit_config, "hf_config": hf_config},
        weight_dir_list=safetensors_list,
        device=get_local_torch_device(),
        hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
        hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
        strict=strict_load,
        cpu_offload=fastvideo_args.dit_cpu_offload,
        pin_cpu_memory=fastvideo_args.pin_cpu_memory,
        fsdp_inference=fastvideo_args.use_fsdp_inference,
        # TODO(will): make these configurable
        default_dtype=default_dtype,
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
        output_dtype=None,
        training_mode=fastvideo_args.training_mode,
        enable_torch_compile=fastvideo_args.enable_torch_compile,
        torch_compile_kwargs=fastvideo_args.torch_compile_kwargs,
    )

    total_params = sum(p.numel() for p in model.parameters())
    logger.info("Loaded model with %.2fB parameters", total_params / 1e9)

    assert next(model.parameters()).dtype == default_dtype, (
        "Model dtype does not match default dtype"
    )

    model = model.eval()

    if fastvideo_args.inference_mode and fastvideo_args.dit_layerwise_offload:
        # Check if model has nn.ModuleList for layerwise offload compatibility
        has_module_list = any(
            isinstance(m, nn.ModuleList) for m in model.children()
        )
        if has_module_list:
            enable_layerwise_offload(model)
        else:
            logger.warning(
                "Layerwise offload requested but model %s does not have "
                "nn.ModuleList structure. Skipping layerwise offload.",
                cls_name
            )
    return model
fastvideo.models.loader.component_loader.UpsamplerLoader 
UpsamplerLoader(device=None)

Bases: ComponentLoader

Loader for upsamplers (incl. LTX-2 spatial/temporal upsamplers).

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device
Methods:
fastvideo.models.loader.component_loader.UpsamplerLoader.load 
load(model_path: str, fastvideo_args: FastVideoArgs)

Load the upsampler based on the model path, and inference args.

Source code in fastvideo/models/loader/component_loader.py 
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the upsampler based on the model path, and inference args."""
    config_dict = get_diffusers_config(model=model_path)
    class_name = config_dict.pop("_class_name", None)

    if class_name is None:
        raise ValueError(
            "Model config does not contain a _class_name attribute. "
            "Only diffusers format is supported.")

    # The base PipelineConfig declares ``upsampler_config`` as a
    # single ``UpsamplerConfig`` instance, but Hunyuan15 narrows it
    # to a tuple of two configs (one per SR target). We only treat
    # the attribute as a multi-config when it actually is one;
    # otherwise the LTX-2 branch below handles the single-class
    # path that takes the diffusers config dict directly.
    upsampler_config_attr = getattr(fastvideo_args.pipeline_config,
                                    "upsampler_config", None)
    if isinstance(upsampler_config_attr, list | tuple):
        try:
            upsampler_cfg = deepcopy(upsampler_config_attr[0])
            upsampler_cfg.update_model_config(config_dict)
        except Exception:
            upsampler_cfg = deepcopy(upsampler_config_attr[1])
            upsampler_cfg.update_model_config(config_dict)
    elif class_name == "LTX2LatentUpsampler":
        # LTX-2 pipeline_config does not declare upsampler_config; the
        # `LTX2LatentUpsampler` wrapper takes the raw diffusers config
        # dict directly via LatentUpsamplerConfigurator.
        upsampler_cfg = deepcopy(config_dict)
    else:
        raise AttributeError(
            "pipeline_config.upsampler_config is missing; cannot build "
            f"upsampler config for class {class_name}")

    model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
    model = model_cls(upsampler_cfg)

    target_device = get_local_torch_device()
    upsampler_precision = getattr(fastvideo_args.pipeline_config,
                                  "upsampler_precision", "bf16")
    model = model.to(target_device,
                     dtype=PRECISION_TO_TYPE[upsampler_precision])

    safetensors_list = glob.glob(
        os.path.join(str(model_path), "*.safetensors"))
    if not safetensors_list:
        raise ValueError(f"No safetensors files found in {model_path}")

    if len(safetensors_list) == 1:
        loaded = safetensors_load_file(safetensors_list[0])
    else:
        loaded = {}
        for sf_file in safetensors_list:
            loaded.update(safetensors_load_file(sf_file))

    # The LTX-2 latent upsampler wrapper exposes the actual conv
    # stack at ``self.model``; checkpoint state_dicts may be saved
    # without the ``model.`` prefix when the inner module was
    # serialised directly. Strip / forward as needed so both layouts
    # load cleanly.
    target_module = getattr(model, "model", model)
    if loaded and all(k.startswith("model.") for k in loaded):
        stripped = {k[len("model."):]: v for k, v in loaded.items()}
        target_module.load_state_dict(stripped, strict=True)
    else:
        target_module.load_state_dict(loaded, strict=True)

    return model.eval()
fastvideo.models.loader.component_loader.VAELoader 
VAELoader(device=None)

Bases: ComponentLoader

Loader for VAE.

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device
Methods:
fastvideo.models.loader.component_loader.VAELoader.load 
load(model_path: str, fastvideo_args: FastVideoArgs)

Load the VAE based on the model path, and inference args.

Source code in fastvideo/models/loader/component_loader.py 
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the VAE based on the model path, and inference args."""
    config = get_diffusers_config(model=model_path)
    class_name = config.pop("_class_name")
    config.pop("_name_or_path", None)
    assert class_name is not None, (
        "Model config does not contain a _class_name attribute. Only diffusers format is supported."
    )
    fastvideo_args.model_paths["vae"] = model_path

    from fastvideo.platforms import current_platform

    if fastvideo_args.vae_cpu_offload:
        target_device = (
            torch.device("mps")
            if current_platform.is_mps()
            else torch.device("cpu")
        )
    else:
        target_device = get_local_torch_device()

    with set_default_torch_dtype(
        PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
        if fastvideo_args.pipeline_config.vae_precision
        else torch.bfloat16
    ):
        pipeline_name = fastvideo_args.pipeline_config.__class__.__name__
        is_gen3c = pipeline_name.startswith("Gen3C")
        is_cosmos25 = pipeline_name == "Cosmos25Config"

        # GEN3C: prefer tokenizer-backed VAE checkpoint when available.
        # This aligns latent conditioning with the GEN3C temporal contract.
        if is_gen3c and class_name in (
                "AutoencoderKLWan", "AutoencoderKLGen3CTokenizer"):
            from fastvideo.models.vaes.gen3c_tokenizer_vae import (
                AutoencoderKLGen3CTokenizer)

            dtype = PRECISION_TO_TYPE[
                fastvideo_args.pipeline_config.vae_precision]
            num_frames = int(
                getattr(fastvideo_args.pipeline_config, "num_frames", 121))
            state_t = int(
                getattr(fastvideo_args.pipeline_config, "state_t", 16))
            if state_t > 1 and num_frames > 1:
                target_temporal = max(1,
                                      (num_frames - 1) // (state_t - 1))
            else:
                target_temporal = 8

            jit_dir = self._find_gen3c_jit_tokenizer_dir(model_path)
            if jit_dir is not None:
                vae = AutoencoderKLGen3CTokenizer.from_jit_tokenizer(
                    jit_dir,
                    device=target_device,
                    dtype=dtype,
                    target_temporal_compression=target_temporal,
                    pixel_chunk_duration=num_frames,
                )
                logger.info(
                    "Loaded GEN3C tokenizer VAE from JIT assets in %s (target temporal compression=%d)",
                    jit_dir,
                    target_temporal,
                )
                return vae.eval()

            tokenizer_ckpt = self._find_gen3c_tokenizer_checkpoint(
                model_path)
            if tokenizer_ckpt is not None:
                vae = AutoencoderKLGen3CTokenizer.from_tokenizer_checkpoint(
                    tokenizer_ckpt,
                    device=target_device,
                    dtype=dtype,
                    target_temporal_compression=target_temporal,
                    pixel_chunk_duration=num_frames,
                )
                logger.info(
                    "Loaded GEN3C tokenizer VAE from %s (target temporal compression=%d)",
                    tokenizer_ckpt,
                    target_temporal,
                )
                return vae.eval()
            logger.warning(
                "GEN3C tokenizer VAE checkpoint not found near %s; falling back to configured class %s.",
                model_path,
                class_name,
            )

        # Cosmos2.5 uses a Wan2.1 VAE stored as `tokenizer.safetensors` under the VAE folder.
        if class_name == "AutoencoderKLWan" and is_cosmos25:
            from fastvideo.models.vaes.cosmos25wanvae import Cosmos25WanVAE

            dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
            vae = Cosmos25WanVAE(device=target_device, dtype=dtype)

            weight_path = os.path.join(model_path, "tokenizer.safetensors")
            if not os.path.exists(weight_path):
                raise FileNotFoundError(
                    f"Missing Cosmos2.5 VAE weights: {weight_path}"
                )
            sd = safetensors_load_file(weight_path)
            vae.load_state_dict(sd, strict=False)
            return vae.eval()

        # LTX-2 uses CausalVideoAutoencoder with nested "vae" config
        if class_name == "CausalVideoAutoencoder" and "vae" in config:
            vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
            vae = vae_cls(config).to(target_device)
            if hasattr(vae, "set_tiling_config"):
                vae_config = fastvideo_args.pipeline_config.vae_config
                vae.set_tiling_config(
                    spatial_tile_size_in_pixels=getattr(
                        vae_config, "ltx2_spatial_tile_size_in_pixels", 512),
                    spatial_tile_overlap_in_pixels=getattr(
                        vae_config, "ltx2_spatial_tile_overlap_in_pixels", 64),
                    temporal_tile_size_in_frames=getattr(
                        vae_config, "ltx2_temporal_tile_size_in_frames", 64),
                    temporal_tile_overlap_in_frames=getattr(
                        vae_config,
                        "ltx2_temporal_tile_overlap_in_frames", 24),
                )
        else:
            config.pop("_class_name", None)
            vae_config = fastvideo_args.pipeline_config.vae_config
            vae_config.update_model_arch(config)
            vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
            vae = vae_cls(vae_config).to(target_device)

    # Find all safetensors files
    safetensors_list = glob.glob(
        os.path.join(str(model_path), "*.safetensors"))
    if not safetensors_list:
        raise ValueError(f"No safetensors files found in {model_path}")
    # Common case: a single `.safetensors` checkpoint file.
    # Some models may be sharded into multiple files; in that case we merge.
    loaded = {}
    for sf_file in safetensors_list:
        loaded.update(safetensors_load_file(sf_file))

    # LTX-2 CausalVideoAutoencoder needs per_channel_statistics remapping
    if class_name == "CausalVideoAutoencoder" and "vae" in config:
        per_channel_prefixes = (
            "per_channel_statistics.",
            "vae.per_channel_statistics.",
        )
        remapped = {}
        for key, tensor in loaded.items():
            remapped[key] = tensor
            for prefix in per_channel_prefixes:
                if key.startswith(prefix):
                    suffix = key[len(prefix):]
                    remapped.setdefault(
                        f"encoder.per_channel_statistics.{suffix}",
                        tensor,
                    )
                    remapped.setdefault(
                        f"decoder.per_channel_statistics.{suffix}",
                        tensor,
                    )
                    break
        loaded = remapped

    # Diffusers-format AutoencoderKL checkpoints should match exactly; load
    # strictly so missing/unexpected keys are surfaced early.
    strict_load = class_name == "AutoencoderKL"
    vae.load_state_dict(loaded, strict=strict_load)
    if (class_name == "AutoencoderKLWan"
            and getattr(vae.config, "use_light_vae", False)
            and target_device.type == "cuda"
            and hasattr(vae, "optimize_memory_format")):
        vae.optimize_memory_format()

    return vae.eval()
fastvideo.models.loader.component_loader.VocoderLoader 
VocoderLoader(device=None)

Bases: ComponentLoader

Loader for LTX-2 vocoder.

Source code in fastvideo/models/loader/component_loader.py 
46
47
def __init__(self, device=None) -> None:
    self.device = device

Functions:

fastvideo.models.loader.fsdp_load

Functions:

fastvideo.models.loader.fsdp_load.load_model_from_full_model_state_dict 
load_model_from_full_model_state_dict(model: FSDPModule | Module, full_sd_iterator: Generator[tuple[str, Tensor], None, None], device: device, param_dtype: dtype, strict: bool = False, cpu_offload: bool = False, param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None, training_mode: bool = True) -> _IncompatibleKeys

Converting full state dict into a sharded state dict and loading it into FSDP model (if training) or normal huggingface model Args: model (Union[FSDPModule, torch.nn.Module]): Model to generate fully qualified names for cpu_state_dict full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs device (torch.device): device used to move full state dict tensors param_dtype (torch.dtype): dtype used to move full state dict tensors strict (bool): flag to check if to load the model in strict mode cpu_offload (bool): flag to check if FSDP offload is enabled param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name training_mode (bool): apply FSDP only for training Returns: NamedTuple with missing_keys and unexpected_keys fields: * missing_keys is a list of str containing the missing keys * unexpected_keys is a list of str containing the unexpected keys

Raises:

Type Description
NotImplementedError

If got FSDP with more than 1D.
Source code in fastvideo/models/loader/fsdp_load.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
def load_model_from_full_model_state_dict(
    model: FSDPModule | torch.nn.Module,
    full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None],
    device: torch.device,
    param_dtype: torch.dtype,
    strict: bool = False,
    cpu_offload: bool = False,
    param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
    training_mode: bool = True,
) -> _IncompatibleKeys:
    """
    Converting full state dict into a sharded state dict
    and loading it into FSDP model (if training) or normal huggingface model
    Args:
        model (Union[FSDPModule, torch.nn.Module]): Model to generate fully qualified names for cpu_state_dict
        full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs
        device (torch.device): device used to move full state dict tensors
        param_dtype (torch.dtype): dtype used to move full state dict tensors
        strict (bool): flag to check if to load the model in strict mode
        cpu_offload (bool): flag to check if FSDP offload is enabled
        param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name
        training_mode (bool): apply FSDP only for training
    Returns:
        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
            * **missing_keys** is a list of str containing the missing keys
            * **unexpected_keys** is a list of str containing the unexpected keys

    Raises:
        NotImplementedError: If got FSDP with more than 1D.
    """
    meta_sd = model.state_dict()
    named_parameters = dict(model.named_parameters())
    sharded_sd = {}
    custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
        full_sd_iterator, param_names_mapping)  # type: ignore
    for target_param_name, full_tensor in custom_param_sd.items():
        meta_sharded_param = meta_sd.get(target_param_name)
        if meta_sharded_param is None:
            # Some checkpoints include extra entries that are not part of the
            # instantiated model's state_dict (e.g. `_extra_state` keys from
            # some FSDP checkpoint formats). These can be safely skipped.
            if (target_param_name.endswith("._extra_state")
                    or target_param_name.endswith("_extra_state")):
                logger.warning(
                    "Skipping non-parameter checkpoint key: %s",
                    target_param_name,
                )
                continue

            # For non-strict loads, treat this as an "unexpected key" and skip it
            # (mirrors torch.nn.Module.load_state_dict(strict=False)).
            if not strict:
                logger.warning(
                    "Skipping unexpected checkpoint key (not present in model): %s",
                    target_param_name,
                )
                continue

            raise ValueError(
                f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect."
            )
        if not hasattr(meta_sharded_param, "device_mesh"):
            full_tensor = full_tensor.to(device=device, dtype=param_dtype)
            target_param = named_parameters.get(target_param_name)
            weight_loader = getattr(target_param, "weight_loader", None)
            # Gated on a shape mismatch: only fused/stacked params with a custom
            # weight_loader (e.g. Qwen3's merged QKV/gate-up) take this path.
            # Existing models whose unsharded params match the checkpoint shape
            # fall through to the original `sharded_tensor = full_tensor` below.
            if target_param is not None and callable(weight_loader) and tuple(target_param.shape) != tuple(
                    full_tensor.shape):
                loaded_param = nn.Parameter(torch.empty(tuple(target_param.shape),
                                                        device=device,
                                                        dtype=param_dtype),
                                            requires_grad=False)
                for attr_name, attr_value in vars(target_param).items():
                    setattr(loaded_param, attr_name, attr_value)
                weight_loader(loaded_param, full_tensor)
                sharded_tensor = loaded_param.data
            else:
                # In cases where parts of the model aren't sharded, some parameters will be plain tensors.
                sharded_tensor = full_tensor
        else:
            full_tensor = full_tensor.to(device=device, dtype=param_dtype)
            sharded_tensor = distribute_tensor(
                full_tensor,
                meta_sharded_param.device_mesh,
                meta_sharded_param.placements,
            )
            if cpu_offload:
                sharded_tensor = sharded_tensor.cpu()
        sharded_sd[target_param_name] = nn.Parameter(sharded_tensor)

    model.reverse_param_names_mapping = reverse_param_names_mapping
    unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys())
    if unused_keys:
        logger.warning("Found unloaded parameters in meta state dict: %s",
                       unused_keys)

    # List of allowed parameter name patterns
    ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress", "proj_l"]  # Can be extended as needed
    for new_param_name in unused_keys:
        if not any(pattern in new_param_name
                   for pattern in ALLOWED_NEW_PARAM_PATTERNS):
            logger.error("Unsupported new parameter: %s. Allowed patterns: %s",
                         new_param_name, ALLOWED_NEW_PARAM_PATTERNS)
            raise ValueError(
                f"New parameter '{new_param_name}' is not supported. "
                f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed."
            )
        meta_sharded_param = meta_sd.get(new_param_name)
        if not hasattr(meta_sharded_param, "device_mesh"):
            # Initialize with zeros
            sharded_tensor = torch.zeros_like(meta_sharded_param,
                                              device=device,
                                              dtype=param_dtype)
        else:
            # Initialize with zeros and distribute
            full_tensor = torch.zeros_like(meta_sharded_param,
                                           device=device,
                                           dtype=param_dtype)
            sharded_tensor = distribute_tensor(
                full_tensor,
                meta_sharded_param.device_mesh,
                meta_sharded_param.placements,
            )
            if cpu_offload:
                sharded_tensor = sharded_tensor.cpu()
        sharded_sd[new_param_name] = nn.Parameter(sharded_tensor)

    # choose `assign=True` since we cannot call `copy_` on meta tensor
    return model.load_state_dict(sharded_sd, strict=strict, assign=True)
fastvideo.models.loader.fsdp_load.maybe_load_fsdp_model 
maybe_load_fsdp_model(model_cls: type[Module], init_params: dict[str, Any], weight_dir_list: list[str], device: device, hsdp_replicate_dim: int, hsdp_shard_dim: int, default_dtype: dtype, param_dtype: dtype, reduce_dtype: dtype, strict: bool = True, cpu_offload: bool = False, fsdp_inference: bool = False, output_dtype: dtype | None = None, training_mode: bool = True, pin_cpu_memory: bool = True, enable_torch_compile: bool = False, torch_compile_kwargs: dict[str, Any] | None = None) -> Module

Load the model with FSDP if is training, else load the model without FSDP.

Source code in fastvideo/models/loader/fsdp_load.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
def maybe_load_fsdp_model(
    model_cls: type[nn.Module],
    init_params: dict[str, Any],
    weight_dir_list: list[str],
    device: torch.device,
    hsdp_replicate_dim: int,
    hsdp_shard_dim: int,
    default_dtype: torch.dtype,
    param_dtype: torch.dtype,
    reduce_dtype: torch.dtype,
    strict: bool = True,
    cpu_offload: bool = False,
    fsdp_inference: bool = False,
    output_dtype: torch.dtype | None = None,
    training_mode: bool = True,
    pin_cpu_memory: bool = True,
    enable_torch_compile: bool = False,
    torch_compile_kwargs: dict[str, Any] | None = None,
) -> torch.nn.Module:
    """
    Load the model with FSDP if is training, else load the model without FSDP.
    """
    # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
    # manually casting the inputs to the model
    mp_policy = MixedPrecisionPolicy(param_dtype,
                                     reduce_dtype,
                                     output_dtype,
                                     cast_forward_inputs=False)

    set_mixed_precision_policy(
        param_dtype=param_dtype,
        reduce_dtype=reduce_dtype,
        output_dtype=output_dtype,
        mp_policy=mp_policy,
    )

    logger.info("Loading model with default_dtype: %s", default_dtype)
    with set_default_dtype(default_dtype), torch.device("meta"):
        model = model_cls(**init_params)

    # Check if we should use FSDP
    use_fsdp = training_mode or fsdp_inference

    # Disable FSDP for MPS as it's not compatible
    from fastvideo.platforms import current_platform
    if current_platform.is_mps():
        use_fsdp = False
        logger.info("Disabling FSDP for MPS platform as it's not compatible")

    if use_fsdp:
        pin_cpu_memory = pin_cpu_memory and is_pin_memory_available()
        world_size = hsdp_replicate_dim * hsdp_shard_dim
        if not training_mode and not fsdp_inference:
            hsdp_replicate_dim = world_size
            hsdp_shard_dim = 1

        if current_platform.is_npu():
            with torch.device("cpu"):
                device_mesh = init_device_mesh(
                    "npu",
                    # (Replicate(), Shard(dim=0))
                    mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
                    mesh_dim_names=("replicate", "shard"),
                )
        else:
            device_mesh = init_device_mesh(
            "cuda",
            # (Replicate(), Shard(dim=0))
            mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
            mesh_dim_names=("replicate", "shard"),
        )
        shard_model(model,
                    cpu_offload=cpu_offload,
                    reshard_after_forward=True,
                    mp_policy=mp_policy,
                    mesh=device_mesh,
                    fsdp_shard_conditions=model._fsdp_shard_conditions,
                    pin_cpu_memory=pin_cpu_memory)

    weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=True)
    param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
    load_model_from_full_model_state_dict(
        model,
        weight_iterator,
        device,
        default_dtype,
        strict=strict,
        cpu_offload=cpu_offload,
        param_names_mapping=param_names_mapping_fn,
    )
    if hasattr(model, "materialize_non_persistent_buffers"):
        model.materialize_non_persistent_buffers(
            device=device, dtype=default_dtype)
    for n, p in chain(model.named_parameters(), model.named_buffers()):
        if p.is_meta:
            raise RuntimeError(
                f"Unexpected param or buffer {n} on meta device.")
        # Avoid unintended computation graph accumulation during inference
        if isinstance(p, torch.nn.Parameter):
            p.requires_grad = False

    # NVFP4 weight prequantization. We detect by the registered
    # ``quant_method`` on linear layers rather than by a separate flag —
    # construction-time ``NVFP4Config.get_quant_method`` already attached
    # ``NVFP4QuantizeMethod`` to every targeted layer, so the loader's
    # responsibility is just to materialize the per-layer nvfp4 weight /
    # scale buffers from the freshly-loaded bf16 weights. No-op when
    # ``flashinfer`` is not installed (lazy import inside the helper).
    _maybe_convert_model_to_nvfp4(model)

    compile_in_loader = enable_torch_compile and training_mode
    if compile_in_loader:
        compile_kwargs = torch_compile_kwargs or {}
        logger.info("Enabling torch.compile for FSDP training module with kwargs=%s",
                    compile_kwargs)
        model = torch.compile(model, **compile_kwargs)
        logger.info("torch.compile enabled for %s", type(model).__name__)
    return model
fastvideo.models.loader.fsdp_load.set_default_dtype 
set_default_dtype(dtype: dtype) -> Generator[None, None, None]

Context manager to set torch's default dtype.

Parameters:

Name Type Description Default
dtype dtype

The desired default dtype inside the context manager.

 required

Returns:

Name Type Description
ContextManager None

context manager for setting default dtype.
Example

with set_default_dtype(torch.bfloat16): x = torch.tensor([1, 2, 3]) x.dtype torch.bfloat16

Source code in fastvideo/models/loader/fsdp_load.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
@contextlib.contextmanager
def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
    """
    Context manager to set torch's default dtype.

    Args:
        dtype (torch.dtype): The desired default dtype inside the context manager.

    Returns:
        ContextManager: context manager for setting default dtype.

    Example:
        >>> with set_default_dtype(torch.bfloat16):
        >>>     x = torch.tensor([1, 2, 3])
        >>>     x.dtype
        torch.bfloat16


    """
    old_dtype = torch.get_default_dtype()
    torch.set_default_dtype(dtype)
    try:
        yield
    finally:
        torch.set_default_dtype(old_dtype)
fastvideo.models.loader.fsdp_load.shard_model 
shard_model(model, *, cpu_offload: bool, reshard_after_forward: bool = True, mp_policy: MixedPrecisionPolicy | None = MixedPrecisionPolicy(), mesh: DeviceMesh | None = None, fsdp_shard_conditions: list[Callable[[str, Module], bool]] = [], pin_cpu_memory: bool = True) -> None

Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API.

This method will over the model's named modules from the bottom-up and apply shard modules based on whether they meet any of the criteria from shard_conditions.

Parameters:

Name Type Description Default
model TransformerDecoder

Model to shard with FSDP.

 required
shard_conditions List[Callable[[str, Module], bool]]

A list of functions to determine which modules to shard with FSDP. Each function should take module name (relative to root) and the module itself, returning True if FSDP should shard the module and False otherwise. If any of shard_conditions return True for a given module, it will be sharded by FSDP.

 required
cpu_offload bool

If set to True, FSDP will offload parameters, gradients, and optimizer states to CPU.

 required
reshard_after_forward bool

Whether to reshard parameters and buffers after the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy.

 True
mesh Optional[DeviceMesh]

Device mesh to use for FSDP sharding under multiple parallelism. Default to None.

 None
fsdp_shard_conditions List[Callable[[str, Module], bool]]

A list of functions to determine which modules to shard with FSDP.

 []
pin_cpu_memory bool

If set to True, FSDP will pin the CPU memory of the offloaded parameters.

 True

Raises:

Type Description
ValueError

If no layer modules were sharded, indicating that no shard_condition was triggered.
Source code in fastvideo/models/loader/fsdp_load.py 
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
def shard_model(
    model,
    *,
    cpu_offload: bool,
    reshard_after_forward: bool = True,
    mp_policy: MixedPrecisionPolicy | None = MixedPrecisionPolicy(),  # noqa
    mesh: DeviceMesh | None = None,
    fsdp_shard_conditions: list[Callable[[str, nn.Module], bool]] = [],  # noqa
    pin_cpu_memory: bool = True,
) -> None:
    """
    Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API.

    This method will over the model's named modules from the bottom-up and apply shard modules
    based on whether they meet any of the criteria from shard_conditions.

    Args:
        model (TransformerDecoder): Model to shard with FSDP.
        shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine
            which modules to shard with FSDP. Each function should take module name (relative to root)
            and the module itself, returning True if FSDP should shard the module and False otherwise.
            If any of shard_conditions return True for a given module, it will be sharded by FSDP.
        cpu_offload (bool): If set to True, FSDP will offload parameters, gradients, and optimizer
            states to CPU.
        reshard_after_forward (bool): Whether to reshard parameters and buffers after
            the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy
            from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy.
        mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism.
            Default to None.
        fsdp_shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine
            which modules to shard with FSDP.
        pin_cpu_memory (bool): If set to True, FSDP will pin the CPU memory of the offloaded parameters.

    Raises:
        ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered.
    """
    # Check if we should use size-based filtering
    use_size_filtering = os.environ.get("FASTVIDEO_FSDP2_AUTOWRAP", "0") == "1"

    if not fsdp_shard_conditions:
        logger.warning("No FSDP shard conditions provided; nothing will be sharded.")
        return

    fsdp_kwargs = {
        "reshard_after_forward": reshard_after_forward,
        "mesh": mesh,
        "mp_policy": mp_policy,
    }
    if cpu_offload:
        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(
            pin_memory=pin_cpu_memory)

    # iterating in reverse to start with
    # lowest-level modules first
    num_layers_sharded = 0

    if use_size_filtering:
        # Size-based filtering mode
        min_params = int(os.environ.get("FASTVIDEO_FSDP2_MIN_PARAMS", "10000000"))
        logger.info("Using size-based filtering with threshold: %.2fM", min_params / 1e6)

        for n, m in reversed(list(model.named_modules())):
            if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]):
                # Count all parameters
                param_count = sum(p.numel() for p in m.parameters(recurse=True))

                # Skip small modules
                if param_count < min_params:
                    logger.info("Skipping module %s (%.2fM params < %.2fM threshold)", 
                               n, param_count / 1e6, min_params / 1e6)
                    continue

                # Shard this module
                logger.info("Sharding module %s (%.2fM params)", n, param_count / 1e6)
                fully_shard(m, **fsdp_kwargs)
                num_layers_sharded += 1
    else:
        # Shard all modules matching conditions        
        for n, m in reversed(list(model.named_modules())):
            if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]):
                fully_shard(m, **fsdp_kwargs)
                num_layers_sharded += 1

        if num_layers_sharded == 0:
            raise ValueError(
                "No layer modules were sharded. Please check if shard conditions are working as expected."
            )

    # Finally shard the entire model to account for any stragglers
    fully_shard(model, **fsdp_kwargs)

fastvideo.models.loader.utils

Utilities for selecting and loading models.

Functions:

fastvideo.models.loader.utils.get_param_names_mapping 
get_param_names_mapping(mapping_dict: dict[str, str]) -> Callable[[str], tuple[str, Any, Any]]

Creates a mapping function that transforms parameter names using regex patterns.

Parameters:

Name Type Description Default
mapping_dict Dict[str, str]

Dictionary mapping regex patterns to replacement patterns

 required
param_name str

The parameter name to be transformed

 required

Returns:

Type Description
Callable[[str], tuple[str, Any, Any]]

Callable[[str], str]: A function that maps parameter names from source to target format
Source code in fastvideo/models/loader/utils.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def get_param_names_mapping(
        mapping_dict: dict[str, str]) -> Callable[[str], tuple[str, Any, Any]]:
    """
    Creates a mapping function that transforms parameter names using regex patterns.

    Args:
        mapping_dict (Dict[str, str]): Dictionary mapping regex patterns to replacement patterns
        param_name (str): The parameter name to be transformed

    Returns:
        Callable[[str], str]: A function that maps parameter names from source to target format
    """

    def mapping_fn(name: str) -> tuple[str, Any, Any]:
        # Try to match and transform the name using the regex patterns in mapping_dict
        for pattern, replacement in mapping_dict.items():
            match = re.match(pattern, name)
            if match:
                merge_index = None
                total_splitted_params = None
                if isinstance(replacement, tuple):
                    merge_index = replacement[1]
                    total_splitted_params = replacement[2]
                    replacement = replacement[0]
                name = re.sub(pattern, replacement, name)
                return name, merge_index, total_splitted_params

        # If no pattern matches, return the original name
        return name, None, None

    return mapping_fn
fastvideo.models.loader.utils.hf_to_custom_state_dict 
hf_to_custom_state_dict(hf_param_sd: dict[str, Tensor] | Iterator[tuple[str, Tensor]], param_names_mapping: Callable[[str], tuple[str, Any, Any]]) -> tuple[dict[str, Tensor], dict[str, tuple[str, Any, Any]]]

Converts a Hugging Face parameter state dictionary to a custom parameter state dictionary.

Parameters:

Name Type Description Default
hf_param_sd Dict[str, Tensor]

The Hugging Face parameter state dictionary

 required
param_names_mapping Callable[[str], tuple[str, Any, Any]]

A function that maps parameter names from source to target format

 required

Returns:

Name Type Description
custom_param_sd Dict[str, Tensor]

The custom formatted parameter state dict
reverse_param_names_mapping Dict[str, Tuple[str, Any, Any]]

Maps back from custom to hf
Source code in fastvideo/models/loader/utils.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def hf_to_custom_state_dict(
    hf_param_sd: dict[str, torch.Tensor] | Iterator[tuple[str, torch.Tensor]],
    param_names_mapping: Callable[[str], tuple[str, Any, Any]]
) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]:
    """
    Converts a Hugging Face parameter state dictionary to a custom parameter state dictionary.

    Args:
        hf_param_sd (Dict[str, torch.Tensor]): The Hugging Face parameter state dictionary
        param_names_mapping (Callable[[str], tuple[str, Any, Any]]): A function that maps parameter names from source to target format

    Returns:
        custom_param_sd (Dict[str, torch.Tensor]): The custom formatted parameter state dict
        reverse_param_names_mapping (Dict[str, Tuple[str, Any, Any]]): Maps back from custom to hf
    """
    custom_param_sd = {}
    to_merge_params = defaultdict(dict)  # type: ignore
    reverse_param_names_mapping = {}
    if isinstance(hf_param_sd, dict):
        hf_param_sd = hf_param_sd.items()  # type: ignore
    for source_param_name, full_tensor in hf_param_sd:  # type: ignore
        target_param_name, merge_index, num_params_to_merge = param_names_mapping(
            source_param_name)
        reverse_param_names_mapping[target_param_name] = (source_param_name,
                                                          merge_index,
                                                          num_params_to_merge)
        if merge_index is not None:
            to_merge_params[target_param_name][merge_index] = full_tensor
            if len(to_merge_params[target_param_name]) == num_params_to_merge:
                # cat at output dim according to the merge_index order
                sorted_tensors = [
                    to_merge_params[target_param_name][i]
                    for i in range(num_params_to_merge)
                ]
                full_tensor = torch.cat(sorted_tensors, dim=0)
                del to_merge_params[target_param_name]
            else:
                continue
        custom_param_sd[target_param_name] = full_tensor
    return custom_param_sd, reverse_param_names_mapping
fastvideo.models.loader.utils.set_default_torch_dtype 
set_default_torch_dtype(dtype: dtype)

Sets the default torch dtype to the given dtype.

Source code in fastvideo/models/loader/utils.py 
16
17
18
19
20
21
22
@contextlib.contextmanager
def set_default_torch_dtype(dtype: torch.dtype):
    """Sets the default torch dtype to the given dtype."""
    old_dtype = torch.get_default_dtype()
    torch.set_default_dtype(dtype)
    yield
    torch.set_default_dtype(old_dtype)

fastvideo.models.loader.weight_utils

Utilities for downloading and initializing model weights.

Functions:

fastvideo.models.loader.weight_utils.default_weight_loader 
default_weight_loader(param: Tensor, loaded_weight: Tensor) -> None

Default weight loader.

Source code in fastvideo/models/loader/weight_utils.py 
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
def default_weight_loader(param: torch.Tensor,
                          loaded_weight: torch.Tensor) -> None:
    """Default weight loader."""
    try:
        if param.numel() == 1 and loaded_weight.numel() == 1:
            # Sometimes scalar values aren't considered tensors with shapes
            # so if both param and loaded_weight are a scalar,
            # "broadcast" instead of copy
            param.data.fill_(loaded_weight.item())
        else:
            assert param.size() == loaded_weight.size(), (
                f"Attempted to load weight ({loaded_weight.size()}) "
                f"into parameter ({param.size()})")

            param.data.copy_(loaded_weight)
    except Exception:
        # NOTE: This exception is added for the purpose of setting breakpoint to
        # debug weight loading issues.
        raise
fastvideo.models.loader.weight_utils.enable_hf_transfer 
enable_hf_transfer() -> None

automatically activates hf_transfer

Source code in fastvideo/models/loader/weight_utils.py 
46
47
48
49
50
51
52
53
54
55
def enable_hf_transfer() -> None:
    """automatically activates hf_transfer
    """
    if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ:
        try:
            # enable hf hub transfer if available
            import hf_transfer  # type: ignore # noqa
            huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER = True
        except ImportError:
            pass
fastvideo.models.loader.weight_utils.filter_files_not_needed_for_inference 
filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]

Exclude files that are not needed for inference.

See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233

Source code in fastvideo/models/loader/weight_utils.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def filter_files_not_needed_for_inference(
        hf_weights_files: list[str]) -> list[str]:
    """
    Exclude files that are not needed for inference.

    See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233
    """
    blacklist = [
        "training_args.bin",
        "optimizer.bin",
        "optimizer.pt",
        "scheduler.pt",
        "scaler.pt",
    ]
    hf_weights_files = [
        f for f in hf_weights_files
        if not any(f.endswith(x) for x in blacklist)
    ]
    return hf_weights_files
fastvideo.models.loader.weight_utils.maybe_remap_kv_scale_name 
maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None

Remap the name of FP8 k/v_scale parameters.

This function handles the remapping of FP8 k/v_scale parameter names. It detects if the given name ends with a suffix and attempts to remap it to the expected name format in the model. If the remapped name is not found in the params_dict, a warning is printed and None is returned.

Parameters:

Name Type Description Default
name str

The original loaded checkpoint parameter name.

 required
params_dict dict

Dictionary containing the model's named parameters.

 required

Returns:

Name Type Description
str str | None

The remapped parameter name if successful, or the original name if no remapping is needed.
None str | None

If the remapped name is not found in params_dict.
Source code in fastvideo/models/loader/weight_utils.py 
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
    """Remap the name of FP8 k/v_scale parameters.

    This function handles the remapping of FP8 k/v_scale parameter names.
    It detects if the given name ends with a suffix and attempts to remap
    it to the expected name format in the model. If the remapped name is not
    found in the params_dict, a warning is printed and None is returned.

    Args:
        name (str): The original loaded checkpoint parameter name.
        params_dict (dict): Dictionary containing the model's named parameters.

    Returns:
        str: The remapped parameter name if successful, or the original name
             if no remapping is needed.
        None: If the remapped name is not found in params_dict.
    """
    if name.endswith(".kv_scale"):
        logger.warning_once(
            "DEPRECATED. Found kv_scale in the checkpoint. "
            "This format is deprecated in favor of separate k_scale and "
            "v_scale tensors and will be removed in a future release. "
            "Functionally, we will remap kv_scale to k_scale and duplicate "
            "k_scale to v_scale")
        # NOTE: we remap the deprecated kv_scale to k_scale
        remapped_name = name.replace(".kv_scale", ".attn.k_scale")
        if remapped_name not in params_dict:
            logger.warning_once(
                f"Found kv_scale in the checkpoint (e.g. {name}), "
                "but not found the expected name in the model "
                f"(e.g. {remapped_name}). kv_scale is "
                "not loaded.")
            return None
        return remapped_name

    possible_scale_names = [".k_scale", ".v_scale"]
    modelopt_scale_names = [
        ".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale"
    ]
    for scale_name in possible_scale_names:
        if name.endswith(scale_name):
            if any(mo_scale_name in name
                   for mo_scale_name in modelopt_scale_names):
                remapped_name = name.replace(
                    f".self_attn.{scale_name[1]}_proj{scale_name}",
                    f".self_attn.attn{scale_name}")
            else:
                remapped_name = name.replace(scale_name, f".attn{scale_name}")
            if remapped_name not in params_dict:
                logger.warning_once(
                    f"Found {scale_name} in the checkpoint (e.g. {name}), "
                    "but not found the expected name in the model "
                    f"(e.g. {remapped_name}). {scale_name} is "
                    "not loaded.")
                return None
            return remapped_name

    # If there were no matches, return the untouched param name
    return name
fastvideo.models.loader.weight_utils.pt_weights_iterator 
pt_weights_iterator(hf_weights_files: list[str], to_cpu: bool = False, broadcast: bool = True) -> Generator[tuple[str, Tensor], None, None]

Iterate over the weights in the model bin/pt files.

Parameters:

Name Type Description Default
hf_weights_files list[str]

List of bin/pt files to load.

 required
to_cpu bool

Whether to load the weights to CPU.

 False
broadcast bool

Accepted for API symmetry. PT weights are loaded through torch.load and do not use the safetensors broadcast path.

 True
Source code in fastvideo/models/loader/weight_utils.py 
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
def pt_weights_iterator(
    hf_weights_files: list[str],
    to_cpu: bool = False,
    broadcast: bool = True
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Iterate over the weights in the model bin/pt files.

    Args:
        hf_weights_files: List of bin/pt files to load.
        to_cpu: Whether to load the weights to CPU.
        broadcast: Accepted for API symmetry. PT weights are loaded through
            torch.load and do not use the safetensors broadcast path.
    """
    node_group = _get_initialized_node_group()
    local_rank = node_group.local_rank if node_group is not None else int(
        os.environ.get("LOCAL_RANK", 0))
    device = str(parallel_state.get_local_torch_device()) if not to_cpu else "cpu"
    enable_tqdm = not torch.distributed.is_initialized() or local_rank == 0
    for bin_file in tqdm(
            hf_weights_files,
            desc="Loading pt checkpoint shards",
            disable=not enable_tqdm,
            bar_format=_BAR_FORMAT,
    ):
        state = torch.load(bin_file, map_location=device, weights_only=True)
        yield from state.items()
        del state
fastvideo.models.loader.weight_utils.resolve_safetensors_files 
resolve_safetensors_files(model_path: str) -> list[str]

Discover safetensors files in a model directory.

Source code in fastvideo/models/loader/weight_utils.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
def resolve_safetensors_files(model_path: str) -> list[str]:
    """Discover safetensors files in a model directory."""
    files = sorted(
        glob.glob(os.path.join(model_path, "*.safetensors")))
    if not files:
        raise FileNotFoundError(
            f"No .safetensors files found in {model_path}")
    index_file = os.path.join(
        model_path, SAFE_WEIGHTS_INDEX_NAME)
    if os.path.exists(index_file):
        files = filter_duplicate_safetensors_files(
            files, model_path, SAFE_WEIGHTS_INDEX_NAME)
    return files
fastvideo.models.loader.weight_utils.safetensors_weights_iterator 
safetensors_weights_iterator(hf_weights_files: list[str], to_cpu: bool = False, broadcast: bool = True, async_broadcast: bool = False) -> Generator[tuple[str, Tensor], None, None]

Iterate over the weights in the model safetensor files. Args: hf_weights_files: List of safetensor files to load. to_cpu: Whether to load the weights to CPU. If False, will load to the GPU device bound to the current process. broadcast: Whether local rank 0 should read GPU weights and broadcast them to the other local ranks. async_broadcast: Whether to overlap loading from disk and broadcasting to other ranks. If True, must iterate over all the weights before use. Only used when broadcast is True and to_cpu is False.

Source code in fastvideo/models/loader/weight_utils.py 
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
def safetensors_weights_iterator(
    hf_weights_files: list[str],
    to_cpu: bool = False,
    broadcast: bool = True,
    async_broadcast: bool = False
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Iterate over the weights in the model safetensor files.
    Args:
        hf_weights_files: List of safetensor files to load.
        to_cpu: Whether to load the weights to CPU. If False, will load to the GPU device bound to the current
            process.
        broadcast: Whether local rank 0 should read GPU weights and broadcast them to the other local ranks.
        async_broadcast: Whether to overlap loading from disk and broadcasting to other ranks. If True,
            must iterate over all the weights before use. Only used when broadcast is True and to_cpu is False.
    """
    node_group = _get_initialized_node_group()
    local_rank = node_group.local_rank if node_group is not None else int(
        os.environ.get("LOCAL_RANK", 0))
    device = str(parallel_state.get_local_torch_device()) if not to_cpu else "cpu"
    enable_tqdm = not torch.distributed.is_initialized() or local_rank == 0
    if to_cpu or not broadcast or node_group is None:
        async_broadcast = False

    handles = []
    for st_file in tqdm(
            hf_weights_files,
            desc="Loading safetensors checkpoint shards",
            disable=not enable_tqdm,
            bar_format=_BAR_FORMAT,
    ):
        with safe_open(st_file, framework="pt", device=device) as f:
            for name in f.keys():  # noqa: SIM118
                if to_cpu:
                    param = f.get_tensor(name)
                elif broadcast and node_group is not None:
                    if local_rank == 0:
                        param = f.get_tensor(name)
                    else:
                        sl = f.get_slice(name)
                        shape = sl.get_shape()
                        dtype = SAFETENSORS_TO_TORCH_DTYPE[sl.get_dtype()]
                        param = torch.empty(shape, device=device, dtype=dtype)
                    # broadcast to local ranks
                    # TODO(Wenxuan): scatter instead of broadcast
                    if node_group.world_size > 1:
                        group = node_group.device_group
                        if async_broadcast:
                            handle = dist.broadcast(param,
                                                    src=dist.get_global_rank(
                                                        group, 0),
                                                    async_op=True,
                                                    group=group)
                            handles.append(handle)
                        else:
                            dist.broadcast(param,
                                           src=dist.get_global_rank(group, 0),
                                           group=group)
                else:
                    param = f.get_tensor(name)
                yield name, param

        if async_broadcast:
            for handle in handles:
                handle.wait()
            handles.clear()

Modules