lora_pipeline ¶

def convert_to_lora_layers(self) -> None:
    """
    Unified method to convert the transformer to a LoRA transformer.
    """
    if self.lora_initialized:
        return
    self.lora_initialized = True
    for (
            transformer_name,
            transformer_module,
    ) in self.trainable_transformer_modules.items():
        if is_lazy_module(transformer_module):

            def _drop_lora_refs(*, _name: str = transformer_name) -> None:
                self.lora_layers.pop(_name, None)
                self.cur_adapter_name = ""
                self.cur_adapter_path = ""

            def _lora_after_load(module: nn.Module, *, _name: str = transformer_name) -> nn.Module:
                self._convert_one_transformer(_name, module)
                self._apply_constructor_adapter()
                return module

            transformer_module.add_release_callback(_drop_lora_refs)
            transformer_module.set_materialize_transform(_lora_after_load)
            continue
        self._convert_one_transformer(transformer_name, transformer_module)

def set_lora_adapter(self,
                     lora_nickname: str,
                     lora_path: str | None = None,
                     strength: float = 1.0,
                     accumulate: bool = False):  # type: ignore
    """
    Load a LoRA adapter into the pipeline and merge it into the transformer.
    Args:
        lora_nickname: The "nick name" of the adapter when referenced in the pipeline.
        lora_path: The path to the adapter, either a local path or a Hugging Face repo id.
        strength: Scale for the low-rank adapter. Hybrid adapters must set this at construction
            so their dense payload receives the same scale.
        accumulate: Add this adapter to an already merged pure low-rank adapter.
    """

    if not math.isfinite(strength):
        raise ValueError(f"LoRA strength must be finite, got {strength}")

    requested_path = lora_path or self.lora_adapter_paths.get(lora_nickname)
    exact_current_adapter = (self.cur_adapter_name == lora_nickname and self.cur_adapter_path == requested_path
                             and self.cur_adapter_strength == strength and not accumulate)
    if exact_current_adapter:
        return

    if not self._setting_constructor_adapter and _has_nvfp4_weights_without_bf16(
            self.trainable_transformer_modules):
        # TODO(David): Restore the BF16 weights and requantize them after an NVFP4 LoRA adapter change.
        raise RuntimeError(
            "Runtime LoRA adapter changes are unsupported after NVFP4 quantization removed the BF16 weights. "
            "Create a new VideoGenerator with the desired LoRA adapter.")

    if not self._setting_constructor_adapter:
        if self._constructor_dense_lora_path is not None:
            raise RuntimeError(
                "The active LoRA contains constructor-time .diff/.set_weight payload. "
                "Changing its adapter or strength at runtime would leave that dense payload stale; "
                "create a new VideoGenerator with ComponentConfig(lora_path=..., lora_strength=...).")
        if requested_path is not None and DenseLoRAPatch.from_adapter(requested_path) is not None:
            raise RuntimeError(
                "Adapters containing .diff/.set_weight payload must be supplied when VideoGenerator is "
                "constructed with ComponentConfig(lora_path=..., lora_strength=...).")

    if lora_nickname not in self.lora_adapters and lora_path is None:
        raise ValueError(f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it.")
    if not self.lora_initialized:
        self.convert_to_lora_layers()
    adapter_updated = False
    rank = dist.get_rank()
    if lora_path is not None and self.lora_adapter_paths.get(lora_nickname) != lora_path:
        self.lora_adapters[lora_nickname] = {}
        lora_local_path = maybe_download_lora(lora_path)
        lora_state_dict = load_file(lora_local_path)

        # Map the hf layer names to our custom layer names
        param_names_mapping_fn = get_param_names_mapping(self.modules["transformer"].param_names_mapping)
        lora_param_names_mapping_fn = get_param_names_mapping(self.modules["transformer"].lora_param_names_mapping)

        # Extract alpha values and weights in a single pass
        to_merge_params: defaultdict[Hashable, dict[Any, Any]] = (defaultdict(dict))
        for name, weight in lora_state_dict.items():
            # Extract weights (lora_A, lora_B, and lora_alpha)
            normalized = normalize_lora_key(name)
            if normalized is None:
                continue
            name = normalized
            name = name.replace(".weight", "")

            if "lora_alpha" in name:
                # Store alpha with minimal mapping - same processing as lora_A/lora_B
                # but store in lora_adapters with ".lora_alpha" suffix
                layer_name = name.replace(".lora_alpha", "")
                layer_name, _, _ = lora_param_names_mapping_fn(layer_name)
                target_name, _, _ = param_names_mapping_fn(layer_name)
                # Store alpha alongside weights with same target_name base
                alpha_key = target_name + ".lora_alpha"
                self.lora_adapters[lora_nickname][alpha_key] = (weight.item()
                                                                if weight.numel() == 1 else float(weight.mean()))
                continue

            name, _, _ = lora_param_names_mapping_fn(name)
            target_name, merge_index, num_params_to_merge = (param_names_mapping_fn(name))
            # for (in_dim, r) @ (r, out_dim), we only merge (r, out_dim * n) where n is the number of linear layers to fuse
            # see param mapping in HunyuanVideoArchConfig
            if merge_index is not None and "lora_B" in name:
                to_merge_params[target_name][merge_index] = weight
                if len(to_merge_params[target_name]) == num_params_to_merge:
                    # cat at output dim according to the merge_index order
                    sorted_tensors = [to_merge_params[target_name][i] for i in range(num_params_to_merge)]
                    weight = torch.cat(sorted_tensors, dim=1)
                    del to_merge_params[target_name]
                else:
                    continue

            if target_name in self.lora_adapters[lora_nickname]:
                raise ValueError(f"Target name {target_name} already exists in lora_adapters[{lora_nickname}]")
            self.lora_adapters[lora_nickname][target_name] = weight.to(self.device)
        adapter_updated = True
        self.cur_adapter_path = lora_path
        self.lora_adapter_paths[lora_nickname] = lora_path
        logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path)

    if (not adapter_updated and self.cur_adapter_name == lora_nickname and self.cur_adapter_strength == strength
            and not accumulate):
        return
    self.cur_adapter_name = lora_nickname
    self.cur_adapter_strength = strength

    # Merge the new adapter
    adapted_count = 0
    consumed: set[str] = set()
    for (
            transformer_name,
            transformer_lora_layers,
    ) in self.lora_layers.items():
        for (
                module,
                layers,
        ) in transformer_lora_layers.lora_layers_by_block():
            with _get_hook_ctx(module):
                for name, layer in layers.items():
                    lora_A_name = name + ".lora_A"
                    lora_B_name = name + ".lora_B"
                    lora_alpha_name = name + ".lora_alpha"
                    if (lora_A_name in self.lora_adapters[lora_nickname]
                            and lora_B_name in self.lora_adapters[lora_nickname]):
                        # Get alpha value for this layer (defaults to None if not present)
                        lora_A = self.lora_adapters[lora_nickname][lora_A_name]
                        lora_B = self.lora_adapters[lora_nickname][lora_B_name]
                        # Simple lookup - alpha stored with same naming scheme as lora_A/lora_B
                        alpha = self.lora_adapters[lora_nickname].get(lora_alpha_name)
                        try:
                            layer.set_lora_weights(
                                lora_A,
                                lora_B,
                                lora_alpha=alpha,
                                training_mode=self.fastvideo_args.training_mode,
                                lora_path=lora_path,
                                strength=strength,
                                accumulate=accumulate,
                            )
                        except Exception as e:
                            logger.error(
                                "Error setting LoRA weights for layer %s: %s",
                                name,
                                str(e),
                            )
                            raise e
                        adapted_count += 1
                        consumed.update((lora_A_name, lora_B_name, lora_alpha_name))
                    else:
                        if rank == 0:
                            logger.warning(
                                "LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.",
                                lora_path,
                                name,
                            )
                        layer.disable_lora = True
    logger.info(
        "Rank %d: LoRA adapter %s applied to %d layers",
        rank,
        lora_path,
        adapted_count,
    )
    # The loop above reports model layers the adapter has nothing for. This is the
    # other direction -- adapter weights that reached no layer -- which is the
    # quieter failure: the adapter loads, generation runs, and the result is simply
    # a partially-applied model with nothing in the log to say so.
    if rank == 0:
        unmatched = sorted(set(self.lora_adapters[lora_nickname]) - consumed)
        for target in unmatched:
            logger.warning("LoRA key not loaded: %s (adapter %s has no matching layer in the model)", target,
                           lora_path)
        if unmatched:
            logger.warning("LoRA adapter %s: %d weights did not reach a layer", lora_path, len(unmatched))
    _convert_quantized_weights_after_lora_merge(self.trainable_transformer_modules)

def unmerge_lora_weights(self) -> None:
    """Unmerge LoRA weights when the transformer's quantized weights remain valid."""
    if _has_quantized_mxfp8_weights(self.trainable_transformer_modules):
        # TODO(David): Requantize MXFP8 weights after LoRA unmerge before enabling this operation.
        raise RuntimeError(
            "LoRA unmerge is unsupported after MXFP8 weight quantization because the quantized weights still "
            "contain the merged LoRA adapter.")
    if _has_quantized_nvfp4_weights(self.trainable_transformer_modules):
        # TODO(David): Preserve BF16 weights and requantize NVFP4 weights after LoRA unmerge.
        raise RuntimeError(
            "LoRA unmerge is unsupported after NVFP4 weight quantization because the quantized weights would "
            "not reflect the unmerged LoRA adapter.")
    for (
            transformer_name,
            transformer_lora_layers,
    ) in self.lora_layers.items():
        for (
                module,
                layers,
        ) in transformer_lora_layers.lora_layers_by_block():
            with _get_hook_ctx(module):
                for name, layer in layers.items():
                    layer.unmerge_lora_weights()