When using torch.no_grad(), will deepspeed *not* perform memory optimization on some parameters of no_grad?

[ojipadeson]

I froze part of the LLM model and introduced it into my custom model. I found that in this case the memory usage of the first GPU is always much higher than the others. Does it mean that deepspeed will not perform memory optimization for the non-gradient part.

[tjruwase]

@ojipadeson, DeepSpeed does optimize for parameters that have no gradients, i.e., frozen parameters. See examples below.

DeepSpeed/deepspeed/runtime/zero/stage_1_and_2.py

Lines 310 to 312 in ffe0af2

	if param.requires_grad:
	param.grad_accum = None
	trainable_parameters.append(param)

DeepSpeed/deepspeed/runtime/zero/stage3.py

Line 446 in ffe0af2

trainable_params = [p for p in param_group[PARAMS_KEY] if p.requires_grad]

Can you please provide repro steps to help investigate what you are observing? Thanks!