GDS offloader#10258
Conversation
…ching - limited to NVIDIA GPUs only
- need to work with tensorflow and other formats - afaik, almost all models shared now is in torch format - converting types should not be that big of a deal
|
Is this supported on Windows? |
|
@yamatazen I only tested on linux (ubuntu), feel free to test it out on windows and share feedback with us. |
Hi, I tried to set this up on Windows without success. I'm getting And I have a suggestion for a small refactor, to move the if hasattr(args, 'enable_gds') and args.enable_gds:
from comfy.gds_loader import GDSConfig, init_gds
config = GDSConfig( ... )
init_gds(gds_config)I'm rooting for this PR, feels like a nice feature! |
|
@bezo97 can you please check if the GDS settings is enabled on your windows machine or not? Maybe windows requires additional drivers sometimes or maybe you need to enable the SDK, afaik they should be there, but you need to enable them. Be careful, when you are working with SSD on windows now, the last month's update broke almost all SSDs. So I can't suggest you any direct solution but I will definetly try to debug the issue together. 
Also we need these deps: Can you confirm please, that till this point everything is working till this point? And then can you please test with these following CLI commands for me? And yes, I am trying to work based on your review now! ref: |
|
@maifeeulasad I made sure the feature is enabled on my Windows. It turns out that the problem is with |
- works only on linux and nvidia
|
@bezo97 thanks, I have pushed the changes! Care to take a look again? |
|
@maifeeulasad LGTM, hopefully others can test it |
|
ping @Kosinkadink @comfyanonymous I have made the changes. Can we merge this now? |
|
I, am testing on ubuntu 24.04 with cuda toolkit 12.9.1 even tough it's work but I when running I have this message Device: cuda:0 NVIDIA GeForce RTX 3080 : cudaMallocAsync is this normal? I try to find workaround but still the same. note : I can run wan 2.1 but its consume all my memory so I need to create swap space same as my memory |
|
@xuge82 I am extremly sorry that you had to face this, care to run these command and share the output with us please? And after you take a note of these, can you please follow the official installation and try install it please? And then run these same three command I shared here and share their output with us. Finally try running comfy again. Hope this resolves your issues. ref: https://docs.nvidia.com/gpudirect-storage/troubleshooting-guide/index.html |
|
Btw, there is an alternative to GDS https://github.com/fal-ai/flashpack But not sure whether it will works on all platforms or not. (i just find out about it). |
|
I think it supports all platforms, no requirements.txt 😄 |
|
Here is my understanding after exploring their codebase and examples so far. If we integrate it on comfyui it we need to rewrite all pipeline for flashpack. Like, if we have And one more thing, we need convert all models into flashpack format, that's an enormous task. Maybe I missed something, open for further discussions. |
|
I think flashpack is just a tensor loader. model need to be transformed to model.flashpack first,then loaded from their own api,that 's all. |
|
I think it will end up like a 'GGUF loader ' node, using the same way like gguf, convert to gguf first then load gguf 😄.Almost all my models are gguf, i don't think it is a problem |
|
I find that their codes still need ram，but need much less ram,very ram friendly. they really have done a lot of optimizations 使用 np.memmap 直接映射文件（零或低拷贝读取），避免反复的 Python 层缓冲与解析开销。memmap 允许按切片读取连续区块。 预计算合理的 chunk 大小并分块读取（pipeline），目标是把 IO 分成 ~100–200 个 chunk，这样可以更好地并行/流水线磁盘读 + host→device 传输。 使用 pinned host buffers（pin_memory=True）作为 staging buffers。这是关键：只有 pinned memory 才能被 cudaMemcpyAsync（或 Tensor.copy_ with non_blocking=True）实现真正的异步、与 CPU-DRAM 复制重叠。PyTorch 官方也推荐 pin_memory + non_blocking 来提速 Host->GPU 传输。 多 CUDA streams + round-robin staging buffers：用多个流把 host→device 的复制异步化，使得当一个 chunk 在 GPU 上传输时，下一个 chunk 可以同时由磁盘/内核读入到另一个 pinned buffer。代码中 num_streams, num_pipeline_buffers 就是为此服务。 调用 madvise(MADV_WILLNEED|MADV_SEQUENTIAL) 提示内核预读/顺序访问，从而降低 page fault/随机读延迟，提升顺序读取吞吐。 针对 bfloat16 做位表示优化（用 uint16 存储位模式），避免每个元素额外转换或不必要的数据膨胀。 |
|
@zwukong 但在windows中使用非阻塞式拷贝特性会引发严重的内存管理问题. 已经被pin住的内存占用将无法被python或torch以任何方式回收, 并且也几乎无法复用 (除非把python进程给kill掉). 你会眼睁睁的看着它把内存全都吃光😂 |
|
@maifeeulasad would using fastsafetensors to wrap the GDS handling simplify integration? |
There was a problem hiding this comment.
Pull request overview
This PR adds GPUDirect Storage (GDS) support to ComfyUI, enabling direct SSD-to-GPU model loading to reduce memory bottlenecks when loading large models. The implementation provides new command-line options, a core GDS loader module, and enhanced node types for model loading with GDS capabilities.
- New GDS loader infrastructure with fallback mechanisms for systems without GDS support
- CLI arguments for configuring GDS behavior (enable/disable, chunk sizes, prefetching, statistics)
- Enhanced checkpoint loader nodes with GDS-aware loading and prefetching capabilities
Reviewed changes
Copilot reviewed 7 out of 8 changed files in this pull request and generated 20 comments.
Show a summary per file
| File | Description |
|---|---|
| comfy/gds_loader.py | Core GDS implementation with model loading, prefetching, and statistics tracking |
| comfy_extras/nodes_gds.py | New ComfyUI nodes for GDS-enabled checkpoint loading, prefetching, and statistics display |
| main.py | GDS initialization logic based on CLI arguments |
| comfy/utils.py | Integration point for GDS in the standard model loading pipeline |
| comfy/cli_args.py | New CLI argument group for GDS configuration options |
| nodes.py | Registration of new GDS nodes module |
| requirements.txt | Formatting fix (newline addition) |
| README.md | Documentation for GDS feature and system requirements |
Comments suppressed due to low confidence (6)
comfy/gds_loader.py:17
- Import of 'Union' is not used.
from typing import Optional, Dict, Any, Union
comfy/gds_loader.py:18
- Import of 'Path' is not used.
from pathlib import Path
comfy/gds_loader.py:26
- Import of 'cuda_runtime' is not used.
import cupy.cuda.runtime as cuda_runtime
comfy/gds_loader.py:33
- Import of 'cudf' is not used.
import cudf # RAPIDS for GPU dataframes
comfy_extras/nodes_gds.py:10
- Import of 'asyncio' is not used.
import asyncio
comfy_extras/nodes_gds.py:17
- Import of 'IO' is not used.
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| try: | ||
| import cudf # RAPIDS for GPU dataframes | ||
| RAPIDS_AVAILABLE = True | ||
| except ImportError: | ||
| RAPIDS_AVAILABLE = False |
There was a problem hiding this comment.
The cudf import is declared but the variable RAPIDS_AVAILABLE is never used anywhere in the code. If RAPIDS/cuDF is not actually required for GDS functionality, consider removing this unused import and flag.
| try: | |
| import cudf # RAPIDS for GPU dataframes | |
| RAPIDS_AVAILABLE = True | |
| except ImportError: | |
| RAPIDS_AVAILABLE = False |
|
|
||
| try: | ||
| import cupy | ||
| import cupy.cuda.runtime as cuda_runtime |
There was a problem hiding this comment.
The cuda_runtime import is unused (line 26). If you need CUDA runtime functions, use them explicitly; otherwise, remove this import.
| import cupy.cuda.runtime as cuda_runtime |
| with open(file_path, 'rb') as f: | ||
| file_size = self._get_file_size(file_path) | ||
|
|
||
| # Choose appropriate buffer or allocate new one | ||
| buffer_size_mb = min(256, max(64, file_size // (1024 * 1024))) | ||
|
|
||
| if buffer_size_mb in self.pinned_buffers: | ||
| pinned_buffer = self.pinned_buffers[buffer_size_mb] | ||
| else: | ||
| # Allocate temporary pinned buffer | ||
| pinned_buffer = cupy.cuda.alloc_pinned_memory(file_size) | ||
|
|
||
| # Read file to pinned memory | ||
| f.readinto(pinned_buffer) | ||
|
|
||
| # Use torch.load with map_location to specific GPU | ||
| # This will be faster due to pinned memory | ||
| state_dict = torch.load( | ||
| f, | ||
| map_location=f'cuda:{self.device}', | ||
| weights_only=True | ||
| ) |
There was a problem hiding this comment.
The file handle f is used with f.readinto(pinned_buffer) on line 281, but then torch.load(f, ...) is called on line 285-289. After readinto consumes the file content, the file pointer is at EOF, so torch.load will try to load from an empty/exhausted file handle. This will cause the PyTorch loading to fail.
The file should either be reopened or the file pointer should be reset with f.seek(0) before calling torch.load. Alternatively, don't use readinto and just pass the file path to torch.load with the pinned memory buffer approach abandoned for PyTorch files.
| cufile.copy_from_file( | ||
| gpu_ptr, | ||
| mmapped_file[start_offset:end_offset], | ||
| tensor_bytes | ||
| ) |
There was a problem hiding this comment.
The cufile.copy_from_file() function signature and usage appears incorrect. According to cuFile documentation, this function doesn't exist in the public API. The correct approach would be to use cufile.CuFile handle's read() method or similar cuFile I/O operations. This implementation will fail at runtime when attempting to use actual GDS functionality.
| import torch | ||
| import time | ||
| from typing import Optional, Dict, Any, Union | ||
| from pathlib import Path |
There was a problem hiding this comment.
The Path import from pathlib is unused in this file. Consider removing it.
| from pathlib import Path |
| import logging | ||
| import torch | ||
| import time | ||
| from typing import Optional, Dict, Any, Union |
There was a problem hiding this comment.
The Union type import from typing is unused in this file. Consider removing it to clean up imports.
| from typing import Optional, Dict, Any, Union | |
| from typing import Optional, Dict, Any |
| """ | ||
|
|
||
| @classmethod | ||
| def INPUT_TYPES(s) -> InputTypeDict: |
There was a problem hiding this comment.
Class methods or methods of a type deriving from type should have 'cls', rather than 's', as their first parameter.
| def INPUT_TYPES(s) -> InputTypeDict: | |
| def INPUT_TYPES(cls) -> InputTypeDict: |
| """ | ||
|
|
||
| @classmethod | ||
| def INPUT_TYPES(s) -> InputTypeDict: |
There was a problem hiding this comment.
Class methods or methods of a type deriving from type should have 'cls', rather than 's', as their first parameter.
| def INPUT_TYPES(s) -> InputTypeDict: | |
| def INPUT_TYPES(cls) -> InputTypeDict: |
| """ | ||
|
|
||
| @classmethod | ||
| def INPUT_TYPES(s) -> InputTypeDict: |
There was a problem hiding this comment.
Class methods or methods of a type deriving from type should have 'cls', rather than 's', as their first parameter.
| def INPUT_TYPES(s) -> InputTypeDict: | |
| def INPUT_TYPES(cls) -> InputTypeDict: |
| from typing import Optional, Dict, Any | ||
|
|
There was a problem hiding this comment.
Import of 'Optional' is not used.
Import of 'Dict' is not used.
Import of 'Any' is not used.
| from typing import Optional, Dict, Any |
|
@guill perhaps it would be better for you to defer your review until after safetensors/safetensors#676 is resolved as that being merged negates the need for this pull request. If you decide to review anyway please read that PR anyway as it anyway as it is by the same author and as least some of the issues raised by the reviewer I believe also apply to this PR such as no benchmarks, looks like AI slop. |
Test Evidence CheckIf this PR changes user-facing behavior, visual proof (screen recording or screenshot) is required. PRs without applicable visual documentation may not be reviewed until provided. You can add it by:
|
|
Hey there @maifeeulasad, would you care to tell me if you used something else other than Nvidia's RAPIDS packages to provide cupy, cudf, and cufile? None of the packages from RAPIDS imports or mentions cupy.cuda.cufile or cufile.CuFileManager. Cufile especially is not being imported from any of these, maybe it's time for a re-base? Btw here's my pip installs from RAPIDS: 
|
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📜 Recent review details
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage | Docstring coverage is 70.97% which is insufficient. The required threshold is 80.00%. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
✅ Passed checks (4 passed)
| Check name | Status | Explanation |
|---|---|---|
| Title check | ✅ Passed | The title is concise and matches the main change: adding GDS offloading support. |
| Description check | ✅ Passed | The description is clearly related to the PR and describes how to test the new GDS workflow. |
| Linked Issues check | ✅ Passed | The changes implement GDS support for NVIDIA/Linux, add an abstract loader path, and surface it through CLI and nodes as requested. |
| Out of Scope Changes check | ✅ Passed | The added README, CLI flags, loader, and node integrations all support the GDS feature and do not appear unrelated. |
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Comment @coderabbitai help to get the list of available commands.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@comfy_extras/nodes_gds.py`:
- Around line 68-123: The node only records GDS intent but always calls
comfy.sd.load_checkpoint_guess_config, so implement an actual GDS-backed load
path: after calling get_gds_instance() and checking
gds._should_use_gds(ckpt_path), call the GDS loader API (e.g., the gds
instance's load method or a provided comfy.gds_loader function) to load the
checkpoint into "out" instead of comfy.sd.load_checkpoint_guess_config;
otherwise fall back to comfy.sd.load_checkpoint_guess_config as before. Ensure
the chosen device (device / target_device) is passed to the loader or used to
set map_location, propagate prefetch success/errors into load_info, and set
load_info["loader_used"] accordingly; reference prefetch_model_gds,
get_gds_instance, gds._should_use_gds, and comfy.sd.load_checkpoint_guess_config
to locate the change points.
In `@comfy/gds_loader.py`:
- Around line 285-308: The _load_fallback function currently returns the raw
object from torch.load or safetensors, which differs from
comfy.utils.load_torch_file; update _load_fallback (and its safetensors branch)
to normalize checkpoints the same way as load_torch_file: if the loaded object
is a dict with a "state_dict" key, return that value; if it's a dict with
exactly one key whose value is a dict of tensors (singleton dict checkpoint),
unwrap and return that inner dict; otherwise return the object as-is. Keep
existing device handling (map_location / safetensors device) and ensure this
normalization is applied to both the safetensors.safe_open result and torch.load
result so callers of _load_fallback receive the same checkpoint shape as
comfy.utils.load_torch_file.
- Around line 440-444: The fallback currently calls comfy.utils.load_torch_file
from the except block in gds_loader, which re-enters the GDS-aware wrapper;
instead, bypass the comfy.utils wrapper and call a non-GDS loader directly (e.g.
torch.load with map_location=device or the original non-wrapped loader function)
so the fallback does not retry GDS and respects --gds-no-fallback; update the
except handler to import torch (or call the original underlying loader) and
return the result using safe_load/device semantics without invoking
comfy.utils.load_torch_file.
- Around line 267-289: The file is read into pinned memory with
f.readinto(pinned_buffer) which advances the file cursor to EOF and the
pinned_buffer is never given to torch.load, so torch.load(f, ...) reads nothing;
fix by either rewinding the file before calling torch.load (f.seek(0)) or by
wrapping the pinned buffer in a file-like object (e.g., io.BytesIO over the
pinned bytes) and pass that to torch.load; update the load logic around
pinned_buffers, the pinned_buffer allocation, and the torch.load call (still
using map_location=f'cuda:{self.device}' and weights_only=True) so that
torch.load reads from the pinned memory rather than an EOF file.
In `@comfy/utils.py`:
- Around line 123-133: The GDS branch in load_torch_file only runs when a CUDA
`device` is passed; update load_torch_file to trigger the GDS path whenever GDS
is enabled (e.g., via the existing enable_gds() flag or config) rather than
requiring device to be non-None, attempt to import and call
gds_loader.load_torch_file_gds even when device argument is None by resolving a
CUDA device internally (or failing back gracefully), and preserve the existing
handling of return_metadata (return (gds_result, {}) when return_metadata True)
and the current exception fallback to the normal loader.
- Around line 127-130: When return_metadata is True, don't return an empty dict;
instead obtain metadata from the non-GDS loader and return it with the GDS
result. Replace the branch that does "return (gds_result, {})" with a fallback
call to the existing loader (e.g., call the original non-GDS loader used
elsewhere in this module with return_metadata=True and the same
ckpt/safe_load/device arguments) and then return (gds_result, metadata) so
gds_loader.load_torch_file_gds provides tensors while the original loader
supplies safetensors/metadata.
In `@main.py`:
- Around line 214-231: The code incorrectly uses hasattr(args, ...) to detect
flags; argparse always defines these attributes so only args.enable_gds is
honored. Update the condition at the top to check the actual boolean values
(e.g., if not (getattr(args,'enable_gds',False) or
getattr(args,'gds_prefetch',False) or getattr(args,'gds_stats',False)): return)
and move the GDSConfig/gds_init block out of the strict args.enable_gds branch
so it runs when any of the flags are true; construct GDSConfig using the parsed
flag values (enable_gds, gds_prefetch, gds_stats, gds_min_file_size,
gds_chunk_size, gds_streams, gds_no_fallback) and call gds_init(config), and
ensure any stats/exit handler registration tied to gds_stats is performed when
getattr(args,'gds_stats',False) is true.
In `@requirements.txt`:
- Line 36: Remove the duplicated dependency "pydantic-settings~=2.0" from
requirements.txt so only a single entry for pydantic-settings remains; locate
both occurrences of "pydantic-settings~=2.0" and delete the redundant line to
avoid duplicate declarations and potential update drift.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 9c4815b2-d782-451a-b69b-3c143440310a
📒 Files selected for processing (8)
README.mdcomfy/cli_args.pycomfy/gds_loader.pycomfy/utils.pycomfy_extras/nodes_gds.pymain.pynodes.pyrequirements.txt
| # Determine target device | ||
| if target_device == "auto": | ||
| device = None # Let the system decide | ||
| elif target_device == "cpu": | ||
| device = torch.device("cpu") | ||
| else: | ||
| device = torch.device(target_device) | ||
|
|
||
| load_info = { | ||
| "file": ckpt_name, | ||
| "path": ckpt_path, | ||
| "target_device": str(device) if device else "auto", | ||
| "gds_enabled": use_gds, | ||
| "prefetch_used": prefetch | ||
| } | ||
|
|
||
| try: | ||
| # Prefetch if requested | ||
| if prefetch and use_gds: | ||
| try: | ||
| from comfy.gds_loader import prefetch_model_gds | ||
| prefetch_success = prefetch_model_gds(ckpt_path) | ||
| load_info["prefetch_success"] = prefetch_success | ||
| if prefetch_success: | ||
| logging.info(f"Prefetched {ckpt_name} to GPU cache") | ||
| except Exception as e: | ||
| logging.warning(f"Prefetch failed for {ckpt_name}: {e}") | ||
| load_info["prefetch_error"] = str(e) | ||
|
|
||
| # Load checkpoint with potential GDS optimization | ||
| if use_gds and device and device.type == 'cuda': | ||
| try: | ||
| from comfy.gds_loader import get_gds_instance | ||
| gds = get_gds_instance() | ||
|
|
||
| # Check if GDS should be used for this file | ||
| if gds._should_use_gds(ckpt_path): | ||
| load_info["loader_used"] = "GDS" | ||
| logging.info(f"Loading {ckpt_name} with GDS") | ||
| else: | ||
| load_info["loader_used"] = "Standard" | ||
| logging.info(f"Loading {ckpt_name} with standard method (file too small for GDS)") | ||
|
|
||
| except Exception as e: | ||
| logging.warning(f"GDS check failed, using standard loading: {e}") | ||
| load_info["loader_used"] = "Standard (GDS failed)" | ||
| else: | ||
| load_info["loader_used"] = "Standard" | ||
|
|
||
| # Load the actual checkpoint | ||
| out = comfy.sd.load_checkpoint_guess_config( | ||
| ckpt_path, | ||
| output_vae=True, | ||
| output_clip=True, | ||
| embedding_directory=folder_paths.get_folder_paths("embeddings") | ||
| ) |
There was a problem hiding this comment.
The node's GDS controls do not change the actual loading path.
target_device and use_gds are only used for load_info/logging here. The real load still goes through comfy.sd.load_checkpoint_guess_config(...), which does not receive the selected device and therefore stays on the standard loader path. Right now this node reports GDS intent, but it does not actually force a GDS-backed load.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@comfy_extras/nodes_gds.py` around lines 68 - 123, The node only records GDS
intent but always calls comfy.sd.load_checkpoint_guess_config, so implement an
actual GDS-backed load path: after calling get_gds_instance() and checking
gds._should_use_gds(ckpt_path), call the GDS loader API (e.g., the gds
instance's load method or a provided comfy.gds_loader function) to load the
checkpoint into "out" instead of comfy.sd.load_checkpoint_guess_config;
otherwise fall back to comfy.sd.load_checkpoint_guess_config as before. Ensure
the chosen device (device / target_device) is passed to the loader or used to
set map_location, propagate prefetch success/errors into load_info, and set
load_info["loader_used"] accordingly; reference prefetch_model_gds,
get_gds_instance, gds._should_use_gds, and comfy.sd.load_checkpoint_guess_config
to locate the change points.
| # Load to pinned memory first | ||
| with open(file_path, 'rb') as f: | ||
| file_size = self._get_file_size(file_path) | ||
|
|
||
| # Choose appropriate buffer or allocate new one | ||
| buffer_size_mb = min(256, max(64, file_size // (1024 * 1024))) | ||
|
|
||
| if buffer_size_mb in self.pinned_buffers: | ||
| pinned_buffer = self.pinned_buffers[buffer_size_mb] | ||
| else: | ||
| # Allocate temporary pinned buffer | ||
| pinned_buffer = cupy.cuda.alloc_pinned_memory(file_size) | ||
|
|
||
| # Read file to pinned memory | ||
| f.readinto(pinned_buffer) | ||
|
|
||
| # Use torch.load with map_location to specific GPU | ||
| # This will be faster due to pinned memory | ||
| state_dict = torch.load( | ||
| f, | ||
| map_location=f'cuda:{self.device}', | ||
| weights_only=True | ||
| ) |
There was a problem hiding this comment.
The staged .pt/.ckpt path consumes the file before torch.load().
After f.readinto(pinned_buffer), the file cursor is at EOF, so torch.load(f, ...) has nothing left to read. The staged buffer is also never passed to torch.load(), so every non-safetensors GDS load will immediately fall into the error path.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@comfy/gds_loader.py` around lines 267 - 289, The file is read into pinned
memory with f.readinto(pinned_buffer) which advances the file cursor to EOF and
the pinned_buffer is never given to torch.load, so torch.load(f, ...) reads
nothing; fix by either rewinding the file before calling torch.load (f.seek(0))
or by wrapping the pinned buffer in a file-like object (e.g., io.BytesIO over
the pinned bytes) and pass that to torch.load; update the load logic around
pinned_buffers, the pinned_buffer allocation, and the torch.load call (still
using map_location=f'cuda:{self.device}' and weights_only=True) so that
torch.load reads from the pinned memory rather than an EOF file.
| state_dict = torch.load( | ||
| f, | ||
| map_location=f'cuda:{self.device}', | ||
| weights_only=True | ||
| ) | ||
|
|
||
| self.stats['gds_loads'] += 1 | ||
| self.stats['total_bytes_gds'] += file_size | ||
|
|
||
| return state_dict | ||
|
|
||
| except Exception as e: | ||
| logging.error(f"GDS PyTorch loading failed: {e}") | ||
| raise | ||
|
|
||
| def _load_fallback(self, file_path: str) -> Dict[str, torch.Tensor]: | ||
| """Fallback loading method using standard approaches""" | ||
| if file_path.lower().endswith(('.safetensors', '.sft')): | ||
| # Use safetensors with device parameter | ||
| with safetensors.safe_open(file_path, framework="pt", device=f'cuda:{self.device}') as f: | ||
| return {k: f.get_tensor(k) for k in f.keys()} | ||
| else: | ||
| # Standard PyTorch loading | ||
| return torch.load(file_path, map_location=f'cuda:{self.device}', weights_only=True) |
There was a problem hiding this comment.
Match comfy.utils.load_torch_file()'s checkpoint normalization.
This replacement returns the raw torch.load(...) object, but the function it is replacing unwraps "state_dict" and singleton-dict checkpoints before returning. For .ckpt/.pt files that changes the return shape and will break downstream model loading/model detection.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@comfy/gds_loader.py` around lines 285 - 308, The _load_fallback function
currently returns the raw object from torch.load or safetensors, which differs
from comfy.utils.load_torch_file; update _load_fallback (and its safetensors
branch) to normalize checkpoints the same way as load_torch_file: if the loaded
object is a dict with a "state_dict" key, return that value; if it's a dict with
exactly one key whose value is a dict of tensors (singleton dict checkpoint),
unwrap and return that inner dict; otherwise return the object as-is. Keep
existing device handling (map_location / safetensors device) and ensure this
normalization is applied to both the safetensors.safe_open result and torch.load
result so callers of _load_fallback receive the same checkpoint shape as
comfy.utils.load_torch_file.
| except Exception as e: | ||
| logging.error(f"GDS loading failed, falling back to standard method: {e}") | ||
| # Fallback to original method | ||
| import comfy.utils | ||
| return comfy.utils.load_torch_file(ckpt, safe_load=safe_load, device=device) |
There was a problem hiding this comment.
The fallback path re-enters GDS recursively.
On failure this calls comfy.utils.load_torch_file(..., device=device), and the new utils wrapper immediately tries GDS again for CUDA devices. That turns a normal fallback into an infinite retry loop and also makes --gds-no-fallback impossible to honor.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@comfy/gds_loader.py` around lines 440 - 444, The fallback currently calls
comfy.utils.load_torch_file from the except block in gds_loader, which re-enters
the GDS-aware wrapper; instead, bypass the comfy.utils wrapper and call a
non-GDS loader directly (e.g. torch.load with map_location=device or the
original non-wrapped loader function) so the fallback does not retry GDS and
respects --gds-no-fallback; update the except handler to import torch (or call
the original underlying loader) and return the result using safe_load/device
semantics without invoking comfy.utils.load_torch_file.
| # Try GDS loading first if available and device is GPU | ||
| if device is not None and device.type == 'cuda': | ||
| try: | ||
| from . import gds_loader | ||
| gds_result = gds_loader.load_torch_file_gds(ckpt, safe_load=safe_load, device=device) | ||
| if return_metadata: | ||
| # For GDS, we return empty metadata for now (can be enhanced) | ||
| return (gds_result, {}) | ||
| return gds_result | ||
| except Exception as e: | ||
| logging.debug(f"GDS loading failed, using fallback: {e}") |
There was a problem hiding this comment.
--enable-gds never reaches the normal checkpoint loader.
This branch only runs when callers explicitly pass a CUDA device, but the built-in checkpoint flow calls load_torch_file(..., return_metadata=True) without one. So enabling GDS at startup does not affect the main model-loading path yet.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@comfy/utils.py` around lines 123 - 133, The GDS branch in load_torch_file
only runs when a CUDA `device` is passed; update load_torch_file to trigger the
GDS path whenever GDS is enabled (e.g., via the existing enable_gds() flag or
config) rather than requiring device to be non-None, attempt to import and call
gds_loader.load_torch_file_gds even when device argument is None by resolving a
CUDA device internally (or failing back gracefully), and preserve the existing
handling of return_metadata (return (gds_result, {}) when return_metadata True)
and the current exception fallback to the normal loader.
| gds_result = gds_loader.load_torch_file_gds(ckpt, safe_load=safe_load, device=device) | ||
| if return_metadata: | ||
| # For GDS, we return empty metadata for now (can be enhanced) | ||
| return (gds_result, {}) |
There was a problem hiding this comment.
Preserve metadata when return_metadata=True.
Returning (gds_result, {}) changes load_torch_file()'s contract for CUDA loads and strips safetensors metadata from downstream consumers. If the GDS path cannot provide metadata yet, it should fall back to the existing loader when metadata is requested instead of silently returning an empty dict.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@comfy/utils.py` around lines 127 - 130, When return_metadata is True, don't
return an empty dict; instead obtain metadata from the non-GDS loader and return
it with the GDS result. Replace the branch that does "return (gds_result, {})"
with a fallback call to the existing loader (e.g., call the original non-GDS
loader used elsewhere in this module with return_metadata=True and the same
ckpt/safe_load/device arguments) and then return (gds_result, metadata) so
gds_loader.load_torch_file_gds provides tensors while the original loader
supplies safetensors/metadata.
| if not hasattr(args, 'enable_gds') and not hasattr(args, 'gds_prefetch') and not hasattr(args, 'gds_stats'): | ||
| # GDS not explicitly requested, use auto-detection | ||
| return | ||
|
|
||
| if hasattr(args, 'enable_gds') and args.enable_gds: | ||
| from comfy.gds_loader import GDSConfig, init_gds as gds_init | ||
|
|
||
| config = GDSConfig( | ||
| enabled=getattr(args, 'enable_gds', False) or getattr(args, 'gds_prefetch', False), | ||
| min_file_size_mb=getattr(args, 'gds_min_file_size', 100), | ||
| chunk_size_mb=getattr(args, 'gds_chunk_size', 64), | ||
| max_concurrent_streams=getattr(args, 'gds_streams', 4), | ||
| prefetch_enabled=getattr(args, 'gds_prefetch', True), | ||
| fallback_to_cpu=not getattr(args, 'gds_no_fallback', False), | ||
| show_stats=getattr(args, 'gds_stats', False) | ||
| ) | ||
|
|
||
| gds_init(config) |
There was a problem hiding this comment.
Use the parsed flag values instead of hasattr() here.
argparse always defines these fields, so this function only initializes GDS inside the args.enable_gds branch. --gds-prefetch and --gds-stats on their own currently do nothing, and the exit stats handler never gets registered.
Suggested fix
- if not hasattr(args, 'enable_gds') and not hasattr(args, 'gds_prefetch') and not hasattr(args, 'gds_stats'):
+ if not (args.enable_gds or args.gds_prefetch or args.gds_stats):
# GDS not explicitly requested, use auto-detection
return
-
- if hasattr(args, 'enable_gds') and args.enable_gds:
+
+ if args.enable_gds or args.gds_prefetch or args.gds_stats:
from comfy.gds_loader import GDSConfig, init_gds as gds_init📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if not hasattr(args, 'enable_gds') and not hasattr(args, 'gds_prefetch') and not hasattr(args, 'gds_stats'): | |
| # GDS not explicitly requested, use auto-detection | |
| return | |
| if hasattr(args, 'enable_gds') and args.enable_gds: | |
| from comfy.gds_loader import GDSConfig, init_gds as gds_init | |
| config = GDSConfig( | |
| enabled=getattr(args, 'enable_gds', False) or getattr(args, 'gds_prefetch', False), | |
| min_file_size_mb=getattr(args, 'gds_min_file_size', 100), | |
| chunk_size_mb=getattr(args, 'gds_chunk_size', 64), | |
| max_concurrent_streams=getattr(args, 'gds_streams', 4), | |
| prefetch_enabled=getattr(args, 'gds_prefetch', True), | |
| fallback_to_cpu=not getattr(args, 'gds_no_fallback', False), | |
| show_stats=getattr(args, 'gds_stats', False) | |
| ) | |
| gds_init(config) | |
| if not (args.enable_gds or args.gds_prefetch or args.gds_stats): | |
| # GDS not explicitly requested, use auto-detection | |
| return | |
| if args.enable_gds or args.gds_prefetch or args.gds_stats: | |
| from comfy.gds_loader import GDSConfig, init_gds as gds_init | |
| config = GDSConfig( | |
| enabled=getattr(args, 'enable_gds', False) or getattr(args, 'gds_prefetch', False), | |
| min_file_size_mb=getattr(args, 'gds_min_file_size', 100), | |
| chunk_size_mb=getattr(args, 'gds_chunk_size', 64), | |
| max_concurrent_streams=getattr(args, 'gds_streams', 4), | |
| prefetch_enabled=getattr(args, 'gds_prefetch', True), | |
| fallback_to_cpu=not getattr(args, 'gds_no_fallback', False), | |
| show_stats=getattr(args, 'gds_stats', False) | |
| ) | |
| gds_init(config) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@main.py` around lines 214 - 231, The code incorrectly uses hasattr(args, ...)
to detect flags; argparse always defines these attributes so only
args.enable_gds is honored. Update the condition at the top to check the actual
boolean values (e.g., if not (getattr(args,'enable_gds',False) or
getattr(args,'gds_prefetch',False) or getattr(args,'gds_stats',False)): return)
and move the GDSConfig/gds_init block out of the strict args.enable_gds branch
so it runs when any of the flags are true; construct GDSConfig using the parsed
flag values (enable_gds, gds_prefetch, gds_stats, gds_min_file_size,
gds_chunk_size, gds_streams, gds_no_fallback) and call gds_init(config), and
ensure any stats/exit handler registration tied to gds_stats is performed when
getattr(args,'gds_stats',False) is true.
maifeeulasad
requested review from
alexisrolland,
kijai and
rattus128
as code owners
|
🎉 Thank you for your contribution, we really appreciate it! 🎉 Like many open source projects, we require contributors to sign our Contributor License Agreement (CLA). A CLA makes the ownership of contributions explicit, so contributors and the project share a clear understanding of how the code can be used. By signing, you:
CLAs are standard practice across major open source projects including those under the Apache Software Foundation and the Linux Foundation. Ours is based on the Apache Software Foundation's CLA. Most importantly, it would enable us to relicense the project under a more permissive license in the future, giving the project and its community greater flexibility. ✍ To sign, please post a new comment on this PR with exactly the following text: ✍ I have read and agree to the Contributor License Agreement You can retrigger this bot by commenting recheck in this Pull Request. Posted by the CLA Assistant Lite bot. |
closes #10257
for testing it on your system:
I was able to run wan 14b without any effort, generally it throws OOM in my GPU, and nearly hangs my CPU.
open for review, feel free to share your thoughts