From cac5a7179900c47279b43b6544648b8a9b7a5fce Mon Sep 17 00:00:00 2001
From: kangwangamd <100359556+kangwangamd@users.noreply.github.com>
Date: Wed, 15 Jul 2026 11:09:56 +0800
Subject: [PATCH] [diffusion] post_training: run weight update under
 torch.inference_mode()

Model params/buffers are created under torch.inference_mode() (they are
'inference tensors'); the in-place copies in the weight-update path
(param.data.copy_, DTensor _local_tensor.copy_, and the layerwise-offload
CPU-buffer copies) must therefore run inside an inference-mode context.
On torch builds with strict inference-tensor enforcement this otherwise
raises 'Inplace update to inference tensor outside InferenceMode is not
allowed' and makes /update_weights_from_disk return an HTTP error,
failing test_update_weights_from_disk.py. Wrapping _load_weights_into_module
in torch.inference_mode() (matching layerwise_offload's own restore path)
fixes it with no CUDA-path change.
---
 .../runtime/post_training/weights_updater.py  | 47 ++++++++++++-------
 1 file changed, 31 insertions(+), 16 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py
index b5d35136b4d0..7c470d0e4ec5 100644
--- a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py
+++ b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py
@@ -169,23 +169,38 @@ def _load_weights_into_module(module: torch.nn.Module, weights_iter) -> None:
 
     For offloaded modules, updates CPU buffers directly via
     update_cpu_weights(); non-offloaded parameters use in-place copy.
+
+    The in-place copies below (param.data.copy_, DTensor _local_tensor.copy_,
+    and the offload manager's CPU-buffer copies) mutate tensors that were
+    created while the pipeline ran under ``torch.inference_mode()`` and are
+    therefore "inference tensors".  PyTorch forbids in-place mutation of an
+    inference tensor outside an inference-mode context, so wrap the whole
+    update in ``torch.inference_mode()`` (matching the offload manager's own
+    restore path).  Without this the weight-update API raises
+    "Inplace update to inference tensor outside InferenceMode is not allowed"
+    and returns an HTTP error.
     """
-    model_params = dict(module.named_parameters())
-    weights_iter = _iter_module_weight_updates(module, weights_iter, model_params)
-
-    offload_managers: list = []
-    if is_layerwise_offloaded_module(module):
-        offload_managers = [m for m in module.layerwise_offload_managers if m.enabled]
-
-    if offload_managers:
-        weight_dict = dict(weights_iter)
-        offloaded_names: set[str] = set()
-        for manager in offload_managers:
-            offloaded_names.update(manager.update_cpu_weights(weight_dict))
-        remaining = ((n, w) for n, w in weight_dict.items() if n not in offloaded_names)
-        load_weights_into_model(remaining, model_params)
-    else:
-        load_weights_into_model(weights_iter, model_params)
+    with torch.inference_mode():
+        model_params = dict(module.named_parameters())
+        weights_iter = _iter_module_weight_updates(module, weights_iter, model_params)
+
+        offload_managers: list = []
+        if is_layerwise_offloaded_module(module):
+            offload_managers = [
+                m for m in module.layerwise_offload_managers if m.enabled
+            ]
+
+        if offload_managers:
+            weight_dict = dict(weights_iter)
+            offloaded_names: set[str] = set()
+            for manager in offload_managers:
+                offloaded_names.update(manager.update_cpu_weights(weight_dict))
+            remaining = (
+                (n, w) for n, w in weight_dict.items() if n not in offloaded_names
+            )
+            load_weights_into_model(remaining, model_params)
+        else:
+            load_weights_into_model(weights_iter, model_params)
 
 
 def _build_module_weight_name_mapper(module: torch.nn.Module):