Loading .safetensors files requires double memory on DGX Spark

[phaserblast]

Updated: A custom loader node for ComfyUI now available:

https://github.com/phaserblast/ComfyUI-DGXSparkSafetensorsLoader

Your question

Loading a single .safetensors files on DGX Spark causes a problem because of the mmap strategy used by the model loader. This happens even when the --disable-mmap option is used. I have been testing FLUX.2-dev, and cannot load the FP16 model despite the DGX Spark having plenty of RAM for loading both the model and text encoder. It seems mmap is a disaster on DGX Spark due to the coherent memory implementation. The result is the safetensors loader tries to load the model twice: First to "RAM," then a copy to "VRAM," which obviously fails since there is no separate RAM/VRAM and we run out of memory loading large models.

This is also a problem with llama-server and LM Studio. With mmap enabled, llama-server tries to load GGUF models into "RAM" first, then copy to "VRAM." Disabling mmap solves the problem, and models can be loaded directly into memory once without the additional move from "RAM" to "VRAM." A similar workaround with the safetensors model loader would be great, and would save time.

Update:

Everything works perfectly and as expected with the BF16 GGUF version of FLUX.2-dev from here:
https://huggingface.co/city96/FLUX.2-dev-gguf
The model loads without doubling the RAM/VRAM requirement, holding just under 90GB with the text encoder also loaded. So the problem is somewhere with the .safetensors loader.

Update 2:

I figured out a way to prevent the ballooning memory when loading a .safetensors file the normal way. It requires an edit to ComfyUI/comfy/utils.py:

if DISABLE_MMAP:  # TODO: Not sure if this is the best way to bypass the mmap issues
    tensor = tensor.to(device=device, copy=True)

Changing copy=True to copy=False allows the model to load without the machine running out of memory:
tensor = tensor.to(device=device, copy=False) # For DGX Spark
Total memory usage is a bit higher than the GGUF version, but inference runs slightly faster. Both models generate identical output.

Remember to launch comfy with the --disable-mmap option.

Logs

Other

No response

[A-Dalton]

I can confirm the finding on my DGX Spark with the same behavior.

Loading the BF16 Safetensor version causes the full 128 GB of unified memory to be used and after the Swap file is also fully utilized, the ComfyUI (Python) process gets killed.
Using the GGUF version (full BF16 GGFUF) in combination with the BF16 Safetensor mistral_3_small_flux2_bf16.safetensors for text encoding only uses roughly 90 GB of the memory.

I was however still unable to run the BF16 Safetensor for FLUX.2 when setting the copy argument to false for the tensor.to function in line 67 in ~/ComfyUI/comfy/utils.py, even when starting ComfyUI with the --disable-mmap option - as it still runs out of memory.

Using the --cache-none option also solves the issue and I can execute the BF16 Safetensor for FLUX.2 with roughly 75 GB of memory used during sampling.

[phaserblast]

Using the --cache-none option also solves the issue and I can execute the BF16 Safetensor for FLUX.2 with roughly 75 GB of memory used during sampling.

Not using --disable-mmap and instead using --cache-none works, but the model is unloaded after every run so this isn't really a solution. There needs to be a way to load the model the "correct" way, but to leave it in memory for the next run without using up every last byte of memory for no reason as is the case now.

These screenshots give an idea of the memory behavior before/after:

Before using --cache-none:

After using --cache-none:

[phaserblast]

I tested this .safetensors loader:
https://github.com/foundation-model-stack/fastsafetensors

This will load a .safetensors file into VRAM directly from storage, bypassing the CPU->GPU copy. It loads the file in only a few seconds, as opposed to a few minutes with the standard safetensors.safe_open method. Perhaps a node can be written to use this loader instead of the default loader, as this is faster and way more efficient on DGX Spark by loading directly into VRAM using NVIDIA GPUDirect. Unfortunately it later fails if I omit the copy using .detach().clone() (pointless though, since it puts things back to square one), so maybe the tensors are falling out of scope before they're accessed. Maybe somewhere down the line after load_torch_file() ComfyUI expects the tensors to be on the CPU:

from fastsafetensors import fastsafe_open
with fastsafe_open(filenames=[ckpt], nogds=False, device=device, debug_log=True) as f:
				sd = {} # state dictionary
				for k in f.keys():
					sd[k] = f.get_tensor(k)

safetensors.safe_open always loads to the CPU using mmap. I'm not sure how the copy is bypassed using --cache-none, but this still isn't a good solution since the model doesn't stay in VRAM after an inference run and it takes forever to load the model again.

[A-Dalton]

Not using --disable-mmap and instead using --cache-none works, but the model is unloaded after every run so this isn't really a solution.

I only mentioned the --cache-none option, as it is a way to execute FLUX.2 with the full BF16 Safetensor variants of both the text encoder and the diffusion model. But I fully agree that it should not be seen as the solution to the problem, if it can be fixed otherwise.

safetensors.safe_open always loads to the CPU using mmap. I'm not sure how the copy is bypassed using --cache-none, but this still isn't a good solution since the model doesn't stay in VRAM after an inference run and it takes forever to load the model again.

To me it looks like the copy to RAM is not actually being bypassed, but rather the ComfyUI cache already contains cached output at the point where it runs OOM.
With --cache-none this output is never cached and therefore only the memory used by the current node is in memory.

Caching happens in the file ~\execution.py when executing a Node with execute_async with the output of the nodes being cached to one of several cache implementations via self.caches.outputs. And the NullCache implementation (used with the argument --cache-none) simply discards the node output instead of caching anything.

With some added additional logging, I can see that there are already some caches set before ComfyUI crashes because it runs OOM (the Killed line):

got prompt
Setting HierarchicalCache cache for Node 10
Using pytorch attention in VAE
Using pytorch attention in VAE
VAE load device: cuda:0, offload device: cuda:0, dtype: torch.bfloat16
Setting HierarchicalCache cache for Node 10
Setting HierarchicalCache cache for Node 47
Setting HierarchicalCache cache for Node 47
Setting HierarchicalCache cache for Node 48
Setting HierarchicalCache cache for Node 48
Setting HierarchicalCache cache for Node 16
Setting HierarchicalCache cache for Node 16
Setting HierarchicalCache cache for Node 38
Using MixedPrecisionOps for text encoder: 210 quantized layers
Requested to load Flux2TEModel_
loaded completely; 95367431640625005117571072.00 MB usable, 1280.59 MB loaded, full load: True
CLIP/text encoder model load device: cuda:0, offload device: cuda:0, current: cuda:0, dtype: torch.float16
Setting HierarchicalCache cache for Node 38
Setting HierarchicalCache cache for Node 38
Setting HierarchicalCache cache for Node 6
Setting HierarchicalCache cache for Node 6
Setting HierarchicalCache cache for Node 6
Setting HierarchicalCache cache for Node 26
Setting HierarchicalCache cache for Node 26
Setting HierarchicalCache cache for Node 12
model weight dtype torch.bfloat16, manual cast: None
model_type FLUX
Killed

The Node Ids being:

• 6 = CLIP Text Encode
• 10 = Load VAE
• 12 = Load Diffusion Model
• 16 = KSamplerSelect
• 26 = FluxGuidance
• 38 = Load CLIP
• 47 = Empty Flux 2 Latent
• 48 = Flux2Schedular

And when using the FP8 version of FLUX.2 with some additional logging I can see that caches from various nodes are being reused on the second execution:

got prompt
NO cached Node Id 6
NO cached Node Id 8
NO cached Node Id 9
Using cached Node Id 10
Using cached Node Id 12
NO cached Node Id 13
Using cached Node Id 16
NO cached Node Id 22
NO cached Node Id 25
NO cached Node Id 26
Using cached Node Id 38
Using cached Node Id 47
Using cached Node Id 48
loaded completely; 95367431640625005117571072.00 MB usable, 17180.59 MB loaded, full load: True

TLDR;

With my limited knowledge it looks more like a limitation of the Safetensors library rather than something ComfyUI can actually solve with their current approach.

I tested this .safetensors loader:
https://github.com/foundation-model-stack/fastsafetensors

Thanks for checking the fastsafetensors library and trying to get it working in ComfyUI!

I hope this might be something which can actually be used in ComfyUI in the future.
While I could not find a sample implementation of fastsafetensors for image generation (to check the different behavior), there is also a discussion thread in the official DGX Spark forum related to LLMs, which points to fastsafetensors bringing a huge improvement:

• https://forums.developer.nvidia.com/t/very-slow-mmap-on-dgx-spark-that-affects-model-loading-questions-to-nvidia/349886

[phaserblast]

I wrote a custom node that does what we want:
https://github.com/phaserblast/ComfyUI-DGXSparkSafetensorsLoader

I only tested it briefly. Give it a shot and let me know if it's working for you.

[A-Dalton]

Really cool - your custom node works pretty well in my initial testing.
Thanks for digging deeper and actually creating a (although very experimental) node for testing purposes!

In combination with the FP8 Mistral_3_small text encoder, image size of 2000x1504, using sageattention and a batch size of 4, your node is always faster when comparing it to the built-in diffusion model loader combined with the --cache-none argument.

Comparing it to the full BF16 GGUF I get about the same memory footprint (which is a bit higher than what you posted using htop) with the stable 2.9.1 version of PyTorch using CUDA 13.0.
Timewise your loader is much faster compared the GGUF loader upon the first full run (00:41:08 vs. 00:37:03), but for consecutive full runs I get slightly higher numbers (GGUF: 00:33:36|00:33:25|00:33:22 / YOURS:00:36:11|00:35:52|00:35:47).

Seeing how you were able to integrate it, I would love to see an official fastsafetensors support for ComfyUI, maybe similar to how the attention mechanism can be controlled via an argument.

[DomiStyle]

This happens on Strix Halo as well, FLUX.2 fails to load in BF16 because it loads to RAM first and then to VRAM.

The node by @phaserblast solves this. Full model fits comfortably into VRAM now.

[jovan2009]

I was excited to try this custom node with fastsafetensors but it seems that I am not able to compile fastsafetensors in windows. The features seem very interesting but...

[phaserblast]

Seeing how you were able to integrate it, I would love to see an official fastsafetensors support for ComfyUI, maybe similar to how the attention mechanism can be controlled via an argument.

Thanks for testing. And yes, it would be great if this could be integrated into the normal model loader as an option!

[phaserblast]

The node by @phaserblast solves this. Full model fits comfortably into VRAM now.

That's great news! I was wondering if this would work on AMD the same way, and I'm happy to hear it does.

[phaserblast]

I was excited to try this custom node with fastsafetensors but it seems that I am not able to compile fastsafetensors in windows. The features seem very interesting but...

Yeah, fastsafetensors/cpp/ext.cpp looks quite Unix-ey, so it might need to be compiled with MinGW.

[jovan2009]

I was excited to try this custom node with fastsafetensors but it seems that I am not able to compile fastsafetensors in windows. The features seem very interesting but...

Yeah, fastsafetensors/cpp/ext.cpp looks quite Unix-ey, so it might need to be compiled with MinGW.

Forgive my lack of knowledge, are you saying that I might have a chance compiling fastsafetensors with MingGW in Windows? I will try if that's true, I used MSYS2 in the past but I don't have it on my computer currently. I even compiled apps inside termux on my 2015 android phone, very similar experience, but it's been a while since I dabbled in a "Linux like" terminal. But I could refresh my memory if it's worth trying.

Indeed, in ext.cpp are called some header .h files that I found out are Unix specific but I passed that compilation step finding their windows adapted variants (some projects on GitHub that have this exact purpose). But then I hit a wall when MSVC complained about a lot of "unreferenced functions", it was obvious it doesn't understand the code and I have no chance.

[phaserblast]

Forgive my lack of knowledge, are you saying that I might have a chance compiling fastsafetensors with MingGW in Windows?

I'm not sure. ext.cpp includes sys/mman.h, which isn't part of the normal MinGW install as far as I can tell. You might have better luck running under WSL2 assuming all the CUDA stuff works. I haven't tried it.

Also, this: It looks like someone made a mman lib for Windows: https://github.com/alitrack/mman-win32 Apparently, this is a package that can be installed if you're using MSYS2. I wish I could be more help! I'm just not very familiar with Windows development outside of using WSL2.

[jovan2009]

Yeah, that was one of the files I searched for, I found probably the same project that you are referring to but it didn't help with MSVC, obviously, I was too naive. I will look into it more, thanks for your suggestions.