wehub-resource-sync
90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Liger-Kernel RMSNorm wrapper for Sana video training.
|
|
|
|
Provides a drop-in replacement for ``diffusion.model.norms.RMSNorm`` using
|
|
Liger-Kernel's Triton-fused implementation. Enabled by default when
|
|
liger-kernel is installed. Disable via ``SANA_USE_LIGER=0``.
|
|
|
|
The Triton kernel fuses cast → square → mean → rsqrt → scale into a single
|
|
GPU pass, eliminating ~5-6 intermediate tensor allocations per call.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
from typing import Type
|
|
|
|
import torch.nn as nn
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_rmsnorm_class() -> Type[nn.Module]:
|
|
"""Return the RMSNorm class to use, controlled by ``SANA_USE_LIGER`` env var.
|
|
|
|
Liger is **on by default** when the package is installed.
|
|
Set ``SANA_USE_LIGER=0`` to force the original PyTorch implementation.
|
|
|
|
Returns:
|
|
``SanaLigerRMSNorm`` when liger-kernel is available (and not disabled),
|
|
otherwise the original ``diffusion.model.norms.RMSNorm``.
|
|
"""
|
|
# Explicitly disabled
|
|
if os.environ.get("SANA_USE_LIGER", "") in ("0", "false", "False"):
|
|
from diffusion.model.norms import RMSNorm
|
|
|
|
return RMSNorm
|
|
|
|
# Auto-detect: try importing liger-kernel
|
|
try:
|
|
from liger_kernel.transformers import LigerRMSNorm
|
|
except ImportError:
|
|
from diffusion.model.norms import RMSNorm
|
|
|
|
return RMSNorm
|
|
|
|
class SanaLigerRMSNorm(LigerRMSNorm):
|
|
"""Drop-in replacement matching ``RMSNorm(dim, scale_factor, eps, norm_dim)`` signature.
|
|
|
|
Uses ``casting_mode="gemma"`` so that both the normalisation *and*
|
|
the weight multiplication happen in fp32 before casting back — this
|
|
matches the original Sana RMSNorm behaviour exactly.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
scale_factor: float = 1.0,
|
|
eps: float = 1e-6,
|
|
norm_dim: int = -1,
|
|
) -> None:
|
|
if norm_dim != -1:
|
|
raise ValueError(f"LigerRMSNorm only supports norm_dim=-1, got {norm_dim}")
|
|
if scale_factor != 1.0:
|
|
raise ValueError(f"LigerRMSNorm requires scale_factor=1.0, got {scale_factor}")
|
|
super().__init__(
|
|
hidden_size=dim,
|
|
eps=eps,
|
|
offset=0.0,
|
|
casting_mode="gemma",
|
|
init_fn="ones",
|
|
in_place=True,
|
|
)
|
|
|
|
logger.info("Using Liger-Kernel RMSNorm (Triton-fused, casting_mode=gemma)")
|
|
return SanaLigerRMSNorm
|