Z1/2 init: flatten params on device

[ksugama]

This PR addresses #7677 by flattening parameter tensors on the accelerators instead of the CPU during zero stage 1 and 2 initialization. This should alleviate CPU contention, with the caveat that the optimization is only used when there is enough VRAM to allocate a full copy of the parameter buffers.

On 8 x H100s and a Intel Xeon Platinum 8480+, profiling the initialization of DeepSpeed on 32 layers of Qwen3-30B with Z2 gives the following:

Old = ~382s
New = ~130s

---

If necessary, this optimization can be extended to allowed a tiered system that trades off VRAM space with performance, which might look like the following:

if enough VRAM for 2x model_size:
    naive flatten
else if enough VRAM for model_size / N:
    distributed flatten across N devices
else:
    flatten on CPU

The distributed flatten would involve each device flattening a portion of the parameters and performing an all-gather to assemble the full flattened model. See #7677 for original discussion.

[sfc-gh-truwase]

@ksugama thanks for working on this. Looking forward to when it is ready for review.

[ksugama]

@stas00 I believe available_memory() eventually calls into nvidia-smi. Is this the foot gun you were warning about?

If it is, maybe this should be fixed in a different PR since that problem touches more than is related to these changes

[stas00]

Please remind me what issue I flagged that you are referring to?

available_memory is here:

DeepSpeed/accelerator/cuda_accelerator.py

Lines 186 to 194 in a44fb58

	def available_memory(self, device_index=None):
	if pynvml:
	if device_index is None:
	device_index = self.current_device()
	handle = pynvml.nvmlDeviceGetHandleByIndex(self._get_nvml_gpu_id(device_index))
	info = pynvml.nvmlDeviceGetMemoryInfo(handle)
	return info.free
	else:
	return self.total_memory(device_index) - self.memory_allocated(device_index)

[ksugama]

You mentioned in the discussion in #7677 that in order to check device memory, nvml should not be used and PyTorch API should be used instead

[stas00]

thank you for the refresher, Kento. I see my comment wasn't very precise. I was trying to say that it doesn't work on all devices.

We do want to use nvml in the general case when we want to see the total stats - since any gpu could be shared by multiple processes and torch's counters can only see its own process.

BTW, torch.cuda.memory.mem_get_info is an alternative to pynvml stats, but I think it'll fail just the same on gb10.

So I think the solution can be this:

 def available_memory(self, device_index=None): 
     this_process_view_available_mem = self.total_memory(device_index) - self.memory_allocated(device_index)
     if pynvml: 
         if device_index is None: 
             device_index = self.current_device() 
         handle = pynvml.nvmlDeviceGetHandleByIndex(self._get_nvml_gpu_id(device_index)) 
         try:
             # gb10 will fail this call
             info = pynvml.nvmlDeviceGetMemoryInfo(handle) 
             return info.free
         except:
             return this_process_view_available_mem
     else: 
         return this_process_view_available_mem

cc: @sfc-gh-truwase

[ksugama]

@stas00 Hoping this alleviates the issue you were running into

[ksugama]

@sfc-gh-truwase Would you be able to approve the CI workflows whenever you get the chance? I ran them on a lambda instance and they pass on my end

[stas00]

@stas00 Hoping this alleviates the issue you were running into

@ksugama, if it solves the problem described in the OP of #7677 then it's fantastic.

I will try to find time to run this check myself.

[stas00]

edit: my repro isn't measuring the right thing, since I see it's taking forever for Transformers to load the shards when doing multi-gpu. your PR is definitely making things load much faster.

I think there is some other inefficiency in HF Transformers's DS handling which makes the checkpoint shards load much slower on 8 gpus as compared to 1 gpu.

[stas00]

Amazing work, thank you, @ksugama.

So many people are going to have a much better experience with DS w/ huge models.

[ksugama]

Incredible! Glad I can help. Thank you for reviewing

[stas00]

https://x.com/StasBekman/status/2022354880049082658