diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 8ed772d70..d4ae19ee1 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -387,14 +387,42 @@ class LayerNorm(CustomOp): # NOTE(will): Needed to match behavior of diffusers and wan2.1 even while using # FSDP's MixedPrecisionPolicy class FP32LayerNorm(nn.LayerNorm): + def _cached_fp32_param( + self, attr: str, param: torch.Tensor | None, device: torch.device + ) -> torch.Tensor | None: + if param is None: + return None + + # Keep autograd semantics identical to the old path. The diffusion + # runtime enters here for inference, where grad is disabled. + if torch.is_grad_enabled(): + return param.float().to(device=device) + + key = ( + param.data_ptr(), + param._version, + param.device, + device, + param.dtype, + ) + cache = self.__dict__.get(attr) + if cache is not None and cache[0] == key: + return cache[1] + + fp32_param = param.detach().to(device=device, dtype=torch.float32) + self.__dict__[attr] = (key, fp32_param) + return fp32_param + def forward(self, inputs: torch.Tensor) -> torch.Tensor: origin_dtype = inputs.dtype device = inputs.device + weight = self._cached_fp32_param("_weight_fp32_cache", self.weight, device) + bias = self._cached_fp32_param("_bias_fp32_cache", self.bias, device) return F.layer_norm( inputs.float(), self.normalized_shape, - self.weight.float().to(device=device) if self.weight is not None else None, - self.bias.float().to(device=device) if self.bias is not None else None, + weight, + bias, self.eps, ).to(origin_dtype) diff --git a/python/sglang/multimodal_gen/test/unit/test_fp32_layernorm.py b/python/sglang/multimodal_gen/test/unit/test_fp32_layernorm.py new file mode 100644 index 000000000..ed4616e7d --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_fp32_layernorm.py @@ -0,0 +1,71 @@ +import pytest +import torch +import torch.nn.functional as F + +from sglang.multimodal_gen.runtime.layers.layernorm import FP32LayerNorm + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fp32_layernorm_cache_matches_reference(): + norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16) + inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(): + actual = norm(inputs) + expected = F.layer_norm( + inputs.float(), + norm.normalized_shape, + norm.weight.float().to(device=inputs.device), + norm.bias.float().to(device=inputs.device), + norm.eps, + ).to(inputs.dtype) + + torch.testing.assert_close(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fp32_layernorm_cache_reuses_converted_params(): + norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16) + inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(): + norm(inputs) + weight_cache = norm.__dict__["_weight_fp32_cache"] + bias_cache = norm.__dict__["_bias_fp32_cache"] + + norm(inputs) + + assert norm.__dict__["_weight_fp32_cache"][1] is weight_cache[1] + assert norm.__dict__["_bias_fp32_cache"][1] is bias_cache[1] + assert "_weight_fp32_cache" not in norm.state_dict() + assert "_bias_fp32_cache" not in norm.state_dict() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fp32_layernorm_cache_invalidates_on_param_update(): + norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16) + inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(): + norm(inputs) + first_key, first_weight = norm.__dict__["_weight_fp32_cache"] + + norm.weight.add_(1.0) + norm(inputs) + second_key, second_weight = norm.__dict__["_weight_fp32_cache"] + + assert second_key != first_key + assert second_weight is not first_weight + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fp32_layernorm_grad_mode_preserves_autograd_path(): + norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16) + inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + output = norm(inputs).float().sum() + output.backward() + + assert inputs.grad is not None + assert "_weight_fp32_cache" not in norm.__dict__ + assert "_bias_fp32_cache" not in norm.__dict__