GdsFileCopier wait_io alignment fix corrupting weights for gpt-oss-120b on DGX Spark

[bbrowning]

When using fastsafetensors as part of vLLM on a DGX Spark (Blackwell GB 10, SM121 CUDA device), the GdsFileCopier gets used by default. The noalign parameter passed into wait_io is hardcoded to False, which triggers the alignment fixing code in GdsFileCopier. This is doing something to the weights, corrupting model responses of gpt-oss-120b and likely other models.

Turning on debug logging, I see messages like this while the model is loading:

(EngineCore_DP0 pid=609395) wait_io: fix misalignment, src=0xe1e980000000, misaligned_bytes=8, count=0, tmp=0xe1e940000000                                                                                                                                                                       
(EngineCore_DP0 pid=609395) wait_io: fix misalignment, src=0xe1e980000000, misaligned_bytes=8, count=1073741824, tmp=0xe1e940000000                                                                                                                                                              
(EngineCore_DP0 pid=609395) wait_io: fix misalignment, src=0xe1e980000000, misaligned_bytes=8, count=2147483648, tmp=0xe1e940000000                                                                                                                                                              
(EngineCore_DP0 pid=609395) wait_io: fix misalignment, src=0xe1e980000000, misaligned_bytes=8, count=3221225472, tmp=0xe1e940000000                                                                                                                                                              
(EngineCore_DP0 pid=609395) wait_io: fix misalignment, src=0xe1e980000000, misaligned_bytes=8, count=4294967296, tmp=0xe1e940000000 

I'm able to restore expected model responses by hardcoding noalign=True in a local build of fastsafetensors or by hardcoding nogds=True in the vLLM integration of fastsafetensors.

[takeshi-yoshimura]

Good catch! So, did you modify the line
https://github.com/foundation-model-stack/fastsafetensors/blob/main/fastsafetensors/loader.py#L142
to set noalign=True to avoid the problem?
This problem seems to be introduced in the current main branch, but pip installation should not cause this. Is this correct?

[bbrowning]

Good catch! So, did you modify the line https://github.com/foundation-model-stack/fastsafetensors/blob/main/fastsafetensors/loader.py#L142 to set noalign=True to avoid the problem?

Yes, testing locally that's what I modified.

This problem seems to be introduced in the current main branch, but pip installation should not cause this. Is this correct?

pip installation from the latest 0.1.15 still corrupts the gpt-oss-120b weights when loading on a DGX Spark. I can workaround this in 0.1.15 by adjusting the vLLM integration to hardcode nogds=True instead of False at https://github.com/vllm-project/vllm/blob/1bb17ecb396f911beaa26ab0d3926d46154c7155/vllm/model_executor/model_loader/weight_utils.py#L783

Note that I'm just loading the model off disk, without any special GPU-direct storage. It's just that on the Spark, it ends up using GdsFileCopier in vLLM because the code is setup to always try that first, only falling back to nogds if an error is thrown with the gds route. No error gets thrown, so I'm hitting the alignment logic that is corrupting weights upon load even in released versions.

I haven't tracked down exactly how the weights are being corrupted yet. Are there models where the alignment fixing code in the GdsFileCopier's wait_io method is known to work?

[bbrowning]

I can workaround this for now in vLLM integration by always disabled GDS there, since that avoids GdsFileCopier entirely. However, on my DGX Spark testing with gpt-oss-120b, the GdsFileCopier is able to load the entire weights in 10 seconds while the non-gds route takes 28 seconds. Both are far faster than using the regular safetensors library. But, if we could get this working with the GdsFileCopier without corrupting the model's output this would be an even better improvement on Spark devices for model loading times.

[eugr]

Hmm, when I run with the version from PyPi, it seems to work properly and I see this in the logs:

 WARNING 12-04 19:17:26 [weight_utils.py:804] GDS not enabled, setting `nogds=True`

However, I had some weights corruption when I used it with the nightly version of pyTorch cu130 wheels yesterday.

As for gpt-oss-120b, I've also had it output some random garbage a few days ago, but don't remember if it was with fastsafetensors or not. I guess, I can test.

[eugr]

@bbrowning - Yep, still working, with fastsafetensors too. I built my container about 12 hours ago, so unless something changed in the release version between then and now, there is something in your setup that enables GDS by default. BTW, according to NVIDIA, GPU Direct is not supported on Spark, so I'm surprised that you've got it working at all, and it works.

Based on their explanation, once memory is allocated by CUDA, it's not available for DMA access from CPU side, including NVMe and ConnectX "for performance reasons".

[takeshi-yoshimura]

Now I see the issue much more clearly. Thank you @bbrowning @eugr. Let me summarize what I understand.

• weights in safetensors files for gpt-oss-120b does not have offsets to be aligned to CUDA. So, GDS without the alignment fix likely causes the CUDA alignment error on later tensor computes.
  • Note: noalign=False moves tensors at non-aligned addresses to aligned one if we detect the alignment issue.
  • I believe noalign=True worked because of vllm. As far as I know, vllm does not directly use tensors from weight loader but always copies them to new tensors at aligned addresses.
  • This is only happening on GDS. nogds=True uses CPU-side bounce buffers without alignment issues.
• The critical thing here was the fact that DGX Spark does not support GDS due to its coherency problem and we must not use GdsFileCopier on the platform: NVIDIA's post
  • In our case, the coherency problem means that the alignment fix likely accesses GPU memory before finishing GDS data transfer, and thus, model weights randomly became garbage values. I believe vllm with noalign=True luckily avoided the issue with enough cycles for other tasks during initialization.
  • The cuFile initialization process does not ensure falling back to the compat mode although there are critical problems in used hardware.

If my understanding is correct, I will fix the code with the following:

• Check CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_SUPPORTED as guided in the NVIDIA's post and ensure using nogds=True to avoid the coherency problem.

I'll keep noalign=False since we don't have to remove it now.

[takeshi-yoshimura]

@bbrowning
#39 is my fix regarding this issue. Can you confirm fastsafetensors always disables GDS on DGX Spark with warnings when using nogds=False. As you said, we cannot gain the maximum performance due to the lack of GDS support.

[bbrowning]

Sorry, just getting back to this. I can confirm that with the change from #39, I now get this warning on my Spark when run with vLLM:

/home/bbrowning/src/fastsafetensors/fastsafetensors/copier/gds.py:188: UserWarning: GDS is not supported in this platform but nogds is False. use nogds=True

And, gpt-oss-120b model loads with fastsafetensors for me again without any corruption of the model weights. I believe you're right that it was likely trying to access the GPU memory previously before data transfer finished.

Thank you!

[takeshi-yoshimura]

ok. the patch has just been merged. I will update the version with the fix this week. thanks!