fix: make training importable and runnable on native Windows

cosmos_framework/model/generator/moe/native_kernels.py

@@ -4,8 +4,37 @@
 from typing import Literal
 
 import torch
-import triton
-import triton.language as tl
+
+try:
+    import triton
+    import triton.language as tl
+except ImportError:  # pragma: no cover - platform dependent
+    # Triton publishes no official Windows wheel. The fused MoE kernels below are its only
+    # consumers here, but this module is imported *transitively* by the config graph
+    # (qwen3_vl_moe/moe_kernels.py -> here) even for dense, non-MoE models. Stub the symbols
+    # so import succeeds; any attempt to actually launch a kernel then fails loudly instead
+    # of silently degrading.
+    class _TritonUnavailable:
+        """Placeholder for triton / triton.language when triton is absent."""
+
+        def __getattr__(self, _name: str) -> "_TritonUnavailable":
+            return _TritonUnavailable()
+
+        def __call__(self, *args, **kwargs):
+            # Support bare-decorator use: @triton.jit
+            if len(args) == 1 and callable(args[0]) and not kwargs:
+
+                def _kernel_unavailable(*_a, **_k):
+                    raise RuntimeError(
+                        "Triton is unavailable on this platform, so the fused MoE kernels cannot run. "
+                        "Use a dense (non-MoE) backbone or install a Triton build for this platform."
+                    )
+
+                return _kernel_unavailable
+            return _TritonUnavailable()
+
+    triton = _TritonUnavailable()  # type: ignore[assignment]
+    tl = _TritonUnavailable()  # type: ignore[assignment]
 
 # Set the token group alignment size for experts in MoE. This is implemented by
 # padding each expert size to the next multiple of TOKEN_GROUP_ALIGN_SIZE_M.

cosmos_framework/model/tokenizer/evaluation/lpips_cache.py

@@ -5,7 +5,6 @@
 
 from __future__ import annotations
 
-import fcntl
 import os
 import stat
 import tempfile
@@ -16,6 +15,10 @@
 
 import torch
 
+# fcntl is POSIX-only; win_compat.flock is a passthrough on Linux and an
+# msvcrt-based equivalent on Windows.
+from cosmos_framework.utils import win_compat as fcntl
+
 _ModelT = TypeVar("_ModelT")
 _READY_MARKER_PAYLOAD = b"cosmos3-tokenizer-lpips-cache-ready-v1\n"
 _READY_MARKER_READ_ATTEMPTS = 10

cosmos_framework/scripts/_train.py

@@ -32,7 +32,7 @@
 from cosmos_framework.inference.common.init import init_output_dir, is_rank0
 from cosmos_framework.utils.flags import SMOKE
 from cosmos_framework.trainer import ImaginaireTrainer
-from cosmos_framework.utils import log
+from cosmos_framework.utils import log, win_compat
 
 if TYPE_CHECKING:
     from torch.utils.data import DataLoader
@@ -120,8 +120,18 @@ def train(args: Args) -> None:
         effective_output_dir.mkdir(parents=True, exist_ok=True)
         symlink_path = effective_output_dir / "job"
         if symlink_path.is_symlink() or symlink_path.exists():
-            os.remove(symlink_path)
-        os.symlink(job_dir, symlink_path)
+            # A junction/dir link needs rmdir, a symlink-to-file needs remove.
+            try:
+                os.remove(symlink_path)
+            except (IsADirectoryError, PermissionError, OSError):
+                shutil.rmtree(symlink_path, ignore_errors=True)
+        # Creating a symlink on Windows requires Developer Mode or
+        # SeCreateSymbolicLinkPrivilege; neither is guaranteed. link_dir falls
+        # back to a directory junction and then to a pointer file so a
+        # convenience alias never aborts the run.
+        strategy = win_compat.link_dir(job_dir, symlink_path)
+        if strategy != "symlink":
+            log.warning(f"Could not symlink {symlink_path} -> {job_dir}; used {strategy} instead.")
     if torch.distributed.is_initialized():
         torch.distributed.barrier()
 

cosmos_framework/scripts/dataset_utils.py

@@ -4,7 +4,6 @@
 """Shared utilities for Action dataset export and evaluation workflows."""
 
 import base64
-import fcntl
 import json
 import os
 import subprocess
@@ -15,6 +14,10 @@
 from loguru import logger
 from tqdm import tqdm
 
+# fcntl is POSIX-only; win_compat.flock is a passthrough on Linux and an
+# msvcrt-based equivalent on Windows.
+from cosmos_framework.utils import win_compat as fcntl
+
 DEFAULT_GCS_CREDENTIALS = "credentials/gcp_checkpoint.secret"
 GCS_CREDS_ENV_VAR = "COSMOS_GCP_CHECKPOINT_CREDS"
 

cosmos_framework/trainer/__init__.py

@@ -351,7 +351,13 @@ def __init__(self, config):
         # One fetched batch reused for cp_size consecutive steps when CP is enabled.
         self._cp_data_window = ContextParallelDataWindow()
         # Send a TimeoutError if a training step takes over timeout_period seconds.
-        signal.signal(signal.SIGALRM, functools.partial(misc.timeout_handler, config.trainer.timeout_period))  # type: ignore
+        # SIGALRM is POSIX-only; Windows has no alarm(2) equivalent, so the
+        # hang watchdog is simply disabled there rather than aborting init.
+        self._hang_watchdog_enabled = hasattr(signal, "SIGALRM")
+        if self._hang_watchdog_enabled:
+            signal.signal(signal.SIGALRM, functools.partial(misc.timeout_handler, config.trainer.timeout_period))  # type: ignore
+        else:
+            log.warning("signal.SIGALRM is unavailable on this platform; the step-timeout watchdog is disabled.")
 
     def _fetch_data_batch(
         self,
@@ -669,7 +675,8 @@ def train(
                     if self.config.trainer.run_validation and (epoch_validation_due or iter_validation_due):
                         self.validate(model, dataloader_val, iteration=iteration)
                     # This iteration is successful; reset the timeout signal.
-                    signal.alarm(self.config.trainer.timeout_period)
+                    if self._hang_watchdog_enabled:
+                        signal.alarm(self.config.trainer.timeout_period)
                     self.straggler_detector.generate_report(iteration)
                     if torch_profiler:
                         torch_profiler.step()
@@ -743,7 +750,15 @@ def training_step(
                 ):
                     output_batch, loss = model_ddp.training_step(data, iteration)
             self.callbacks.on_after_forward(iteration=iteration)
-            model = model_ddp.module if self.config.trainer.distributed_parallelism == "ddp" else model_ddp
+            # `.module` only exists when the DDP wrapper was actually applied.
+            # parallel_model_wrapper may hand back the bare module (e.g.
+            # world_size==1 on Windows), so unwrap defensively rather than on
+            # the configured mode alone.
+            model = (
+                model_ddp.module
+                if self.config.trainer.distributed_parallelism == "ddp" and hasattr(model_ddp, "module")
+                else model_ddp
+            )
             self.callbacks.on_before_backward(model, loss, iteration=iteration)
             with self.training_timer("backward"):
                 with self.straggler_detector.profile_section(

cosmos_framework/utils/checkpoint_db.py

@@ -263,6 +263,19 @@ def _download(self) -> str:
                 for i, pattern in enumerate(patterns):
                     patterns[i] = os.path.join(self.subdirectory, pattern)
 
+        # `--include *` is a *trivial* filter -- it already means "everything" --
+        # but the hf CLI parses a bare `*` as a positional FILENAME rather than as
+        # the include pattern:
+        #     UserWarning: Ignoring `--include` since filenames have been explicitly set.
+        # Only README.md is then fetched, `blobs/` stays empty, and the CLI still
+        # exits 0 returning a valid-looking snapshot path -- so the caller proceeds
+        # on an essentially empty snapshot and fails much later and misleadingly
+        # (e.g. `TypeError: expected str, bytes or os.PathLike object, not NoneType`
+        # out of the tokenizer's `open(vocab_file)`). Drop the flag when it adds
+        # nothing; a subdirectory-scoped pattern is still passed through.
+        if include == ["*"]:
+            include = []
+
         cmd_args = [
             self.repository,
             "--repo-type",

cosmos_framework/utils/distributed.py

@@ -38,6 +38,12 @@
 
 def _set_gpu_cpu_affinity(device: Device) -> None:
     """Prefer GPU-local CPUs within the process's existing allocation."""
+    # os.sched_getaffinity / os.sched_setaffinity are POSIX-only and do not exist on
+    # Windows. Guard before calling: an AttributeError is not an OSError, so it would
+    # escape the caller's handler and abort every entrypoint before training starts.
+    if not (hasattr(os, "sched_getaffinity") and hasattr(os, "sched_setaffinity")):
+        log.warning("os.sched_getaffinity/sched_setaffinity are unavailable on this platform; skipping CPU affinity.")
+        return
     allocated_cpus = os.sched_getaffinity(0)
     preferred_cpus = allocated_cpus.intersection(device.get_cpu_affinity())
     if preferred_cpus:
@@ -69,7 +75,7 @@ def init(store: dist.Store | None = None, backend: str | None = None) -> int | N
     try:
         device = Device(local_rank)
         _set_gpu_cpu_affinity(device)
-    except (pynvml.NVMLError, OSError) as e:
+    except (pynvml.NVMLError, OSError, AttributeError) as e:
         log.warning(f"Failed to set device affinity: {e}")
     # Set up distributed communication. CPU checkpoint conversion needs Gloo
     # because NCCL cannot synchronize CPU-resident tokenizer or model tensors.
@@ -92,6 +98,12 @@ def init(store: dist.Store | None = None, backend: str | None = None) -> int | N
             )
         if backend is None:
             backend = "nccl" if os.environ.get("COSMOS_DEVICE", "cuda").lower() == "cuda" else "gloo"
+        # NCCL has no Windows build (torch.distributed.is_nccl_available() is False there).
+        # A process group is still mandatory rather than optional: Trainer.__init__ runs an
+        # unconditional dist.barrier(). Fall back to gloo for collectives while compute stays on CUDA.
+        if sys.platform == "win32" and backend == "nccl":
+            log.warning("NCCL is unavailable on Windows; using the gloo backend instead (compute stays on CUDA).")
+            backend = "gloo"
         if store is not None:
             dist.init_process_group(
                 backend=backend,
@@ -101,7 +113,24 @@ def init(store: dist.Store | None = None, backend: str | None = None) -> int | N
                 timeout=timeout_timedelta,
             )
         else:
+            if sys.platform == "win32":
+                # env:// needs these set even for a single-rank job; there is no torchrun-style
+                # launcher on the native Windows path.
+                os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+                os.environ.setdefault("MASTER_PORT", "29500")
+                os.environ.setdefault("RANK", str(local_rank))
+                os.environ.setdefault("WORLD_SIZE", "1")
             dist.init_process_group(backend=backend, init_method="env://", timeout=timeout_timedelta)
+        if sys.platform == "win32":
+            # The Windows gloo build cannot run collectives on CUDA tensors; it faults with an
+            # access violation. Route them via host memory.
+            from cosmos_framework.utils import win_compat
+
+            if win_compat.install_gloo_cuda_collective_shims():
+                log.warning(
+                    "Installed gloo/CUDA collective shims: torch.distributed collectives will "
+                    "stage CUDA tensors through host memory on this platform."
+                )
         log.critical(
             f"Initialized distributed training with local rank {local_rank} using {backend} with timeout {timeout_seconds}",
             rank0_only=False,
@@ -334,6 +363,18 @@ def parallel_model_wrapper(config_ddp: DDPConfig, model: torch.nn.Module) -> tor
     """
     if dist.is_available() and dist.is_initialized():
         local_rank = int(os.getenv("LOCAL_RANK", 0))
+        # On Windows the only available backend is gloo (NCCL has no Windows build). PyTorch's
+        # DDP reducer runs its gradient allreduce on the CUDA-resident grads, and the Windows
+        # gloo build has no CUDA-aware path for that -- the result is an access violation
+        # (0xC0000005) part way through backward, with no Python traceback. At world_size == 1
+        # the DDP wrapper is pure overhead (there is nothing to reduce), so hand back the bare
+        # module. Guarded on win32 so Linux behaviour is untouched.
+        if sys.platform == "win32" and get_world_size() == 1:
+            log.warning(
+                "Skipping the DDP wrapper: world_size==1 on Windows, where the gloo backend cannot "
+                "allreduce CUDA gradients. Gradients are already rank-local, so this is a no-op."
+            )
+            return model
         try:
             from megatron.core import parallel_state
 

cosmos_framework/utils/easy_io/transient_retry.py

@@ -27,7 +27,19 @@
 from botocore.exceptions import (
     ReadTimeoutError as BotocoreReadTimeoutError,
 )
-from multistorageclient.types import RetryableError
+
+try:
+    from multistorageclient.types import RetryableError
+except ImportError:  # pragma: no cover - platform dependent
+
+    class RetryableError(Exception):  # type: ignore[no-redef]
+        """Stand-in when multi-storage-client is unavailable (e.g. on Windows).
+
+        Kept in RETRYABLE_EXCEPTIONS so the tuple shape is unchanged; nothing
+        ever raises it because the MSC backend is not registered.
+        """
+
+
 from urllib3.exceptions import ProtocolError as URLLib3ProtocolError
 from urllib3.exceptions import ReadTimeoutError as URLLib3ReadTimeoutError
 from urllib3.exceptions import SSLError as URLLib3SSLError

cosmos_framework/utils/reasoner/distributed.py

@@ -3,7 +3,6 @@
 
 # Standard library imports
 import array
-import fcntl
 import math
 import os
 import re
@@ -14,6 +13,15 @@
 from functools import partial
 from typing import Dict, Iterable, Optional, Union
 
+# fcntl is POSIX-only. It is used here solely for SIOCGIFADDR/SIOCGIFCONF
+# interface discovery, which is a multi-node Infiniband/Mellanox concern; a
+# single-node Windows run never needs it. Import defensively so the module stays
+# importable, and let the helpers below degrade to "no interfaces found".
+try:
+    import fcntl
+except ImportError:  # pragma: no cover - platform dependent
+    fcntl = None  # type: ignore[assignment]
+
 # Third party imports
 import requests
 import torch
@@ -263,6 +271,9 @@ def get_ip_address(ifname):
     Returns:
         str or None: The IPv4 address as a string if found, else None.
     """
+    if fcntl is None:
+        # No SIOCGIFADDR equivalent on this platform (single-node only).
+        return None
     s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
     try:
         # 0x8915 is SIOCGIFADDR; pack interface name (limited to 15 chars)
@@ -319,6 +330,10 @@ def get_all_ipv4_addresses():
     ip_list = []
     s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 
+    if fcntl is None:
+        # No SIOCGIFCONF equivalent on this platform (single-node only).
+        return ip_list
+
     # Allocate buffer for maximum number of interfaces.
     max_interfaces = 128
     bytes_size = max_interfaces * 32