Fix fp8 stream: bf16 assign then in-place float8 cast + progress

app.py

@@ -141,7 +141,7 @@ MEM_MODE = (
     os.environ.get("NFA_FUN_CN_MEM_MODE") or "model_cpu_offload_and_qfloat8"
 ).strip()
 # Abort CPU load if Fun CN not ready — do not thrash 35+ min again.
-LOAD_DEADLINE_SEC = int(os.environ.get("NFA_FUN_CN_LOAD_DEADLINE_SEC") or "480")
+LOAD_DEADLINE_SEC = int(os.environ.get("NFA_FUN_CN_LOAD_DEADLINE_SEC") or "900")
 
 CONFIG_PATH = APP_DIR / "config" / "flux2_control.yaml"
 MODEL_DIR = Path(os.environ.get("NFA_FLUX2_MOUNT") or "/data/FLUX.2-dev")
@@ -282,18 +282,31 @@ def _fp8_dtype_for_key(key: str) -> torch.dtype:
 
 
 def _set_tensor(model: torch.nn.Module, key: str, tensor: torch.Tensor) -> None:
+    """Materialize on CPU as bf16, then in-place cast to float8 when allowed.
+
+    Direct float8 via set_module_tensor_to_device hung on ZeroGPU CPU; bf16
+    assign + ``param.data = param.data.to(float8)`` is the stable path.
+    """
     from accelerate.utils import set_module_tensor_to_device
 
-    target = _fp8_dtype_for_key(key)
-    # float8 cast must go through a float dtype first on some builds
-    if target == torch.float8_e4m3fn and tensor.dtype not in (
-        torch.float8_e4m3fn,
-        torch.float8_e5m2,
-    ):
-        value = tensor.detach().to(dtype=torch.bfloat16).to(dtype=target)
-    else:
-        value = tensor.detach().to(dtype=target)
-    set_module_tensor_to_device(model, key, device="cpu", value=value, dtype=target)
+    value = tensor.detach().to(dtype=WEIGHT_DTYPE)
+    set_module_tensor_to_device(
+        model, key, device="cpu", value=value, dtype=WEIGHT_DTYPE
+    )
+    del value
+    if _fp8_dtype_for_key(key) != torch.float8_e4m3fn:
+        return
+    # Walk to the Parameter and cast storage in-place (frees bf16 bits).
+    mod: torch.nn.Module = model
+    parts = key.split(".")
+    for p in parts[:-1]:
+        mod = getattr(mod, p)
+    leaf = parts[-1]
+    param = getattr(mod, leaf)
+    if isinstance(param, torch.nn.Parameter):
+        param.data = param.data.to(torch.float8_e4m3fn)
+    elif isinstance(param, torch.Tensor):
+        setattr(mod, leaf, param.to(torch.float8_e4m3fn))
 
 
 def _stream_shards_into_model(
@@ -303,7 +316,7 @@ def _stream_shards_into_model(
     label: str,
 ) -> None:
     """Load one safetensors shard at a time → fp8; never accumulate full bf16."""
-    from safetensors import safe_open
+    from safetensors.torch import load_file
 
     model_sd = model.state_dict()
     loaded = 0
@@ -315,23 +328,24 @@ def _stream_shards_into_model(
             f"rss_max≈{_rss_gb():.1f}GB path={Path(path).name}",
             flush=True,
         )
-        with safe_open(path, framework="pt", device="cpu") as f:
-            for key in f.keys():
-                if key not in model_sd:
-                    skipped += 1
-                    continue
-                tensor = f.get_tensor(key)
-                if tuple(tensor.shape) != tuple(model_sd[key].shape):
-                    print(
-                        f"[nfa-fun-cn] skip size mismatch {key} "
-                        f"{tuple(tensor.shape)} vs {tuple(model_sd[key].shape)}",
-                        flush=True,
-                    )
-                    skipped += 1
-                    continue
-                _set_tensor(model, key, tensor)
-                loaded += 1
-                del tensor
+        sd = load_file(path, device="cpu")
+        n_keys = len(sd)
+        for j, (key, tensor) in enumerate(sd.items()):
+            if key not in model_sd:
+                skipped += 1
+                continue
+            if tuple(tensor.shape) != tuple(model_sd[key].shape):
+                skipped += 1
+                continue
+            _set_tensor(model, key, tensor)
+            loaded += 1
+            if (j + 1) % 200 == 0 or (j + 1) == n_keys:
+                print(
+                    f"[nfa-fun-cn] {label} shard {i+1} keys {j+1}/{n_keys} "
+                    f"loaded={loaded} rss_max≈{_rss_gb():.1f}GB",
+                    flush=True,
+                )
+        del sd
         gc.collect()
     print(
         f"[nfa-fun-cn] {label} stream done loaded={loaded} skipped={skipped} "