Problem when finetuning LLM using opacus

username: Kai9909
cooked: <p>I tried to finetune a LLM model (distlgpt2 in <a href="http://huggingface.co" rel="noopener nofollow ugc">huggingface.co</a>) using opacus. I just do anything as normal but get an unexpected error.</p>
<p>Traceback (most recent call last):<br>
File “/u/nkp2mr/kaic/dp_diffusion_synthesis/LLM/run_llm.py”, line 25, in <br>
model.fit(data_path)<br>
File “/u/nkp2mr/kaic/dp_diffusion_synthesis/./LLM/great.py”, line 192, in fit<br>
self.model = model_trainer.train(great_ds, ‘finetune’)<br>
File “/u/nkp2mr/kaic/dp_diffusion_synthesis/./LLM/great_trainer.py”, line 106, in train<br>
self.optimizer.step()<br>
File “/u/nkp2mr/.local/lib/python3.10/site-packages/torch/optim/lr_scheduler.py”, line 130, in wrapper<br>
return func.<strong>get</strong>(opt, opt.<strong>class</strong>)(*args, **kwargs)<br>
File “/u/nkp2mr/.local/lib/python3.10/site-packages/opacus/optimizers/optimizer.py”, line 513, in step<br>
if self.pre_step():<br>
File “/u/nkp2mr/.local/lib/python3.10/site-packages/opacus/optimizers/optimizer.py”, line 494, in pre_step<br>
self.clip_and_accumulate()<br>
File “/u/nkp2mr/.local/lib/python3.10/site-packages/opacus/optimizers/optimizer.py”, line 404, in clip_and_accumulate<br>
per_sample_norms = torch.stack(per_param_norms, dim=1).norm(2, dim=1)<br>
RuntimeError: stack expects each tensor to be equal size, but got [13] at entry 0 and [1] at entry 1</p>
<p>I print the model structure:<br>
GPT2LMHeadModel(<br>
(transformer): GPT2Model(<br>
(wte): Embedding(50257, 768)<br>
(wpe): Embedding(1024, 768)<br>
(drop): Dropout(p=0.1, inplace=False)<br>
(h): ModuleList(<br>
(0-5): 6 x GPT2Block(<br>
(ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)<br>
(attn): GPT2SdpaAttention(<br>
(c_attn): Conv1D()<br>
(c_proj): Conv1D()<br>
(attn_dropout): Dropout(p=0.1, inplace=False)<br>
(resid_dropout): Dropout(p=0.1, inplace=False)<br>
)<br>
(ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)<br>
(mlp): GPT2MLP(<br>
(c_fc): Conv1D()<br>
(c_proj): Conv1D()<br>
(act): NewGELUActivation()<br>
(dropout): Dropout(p=0.1, inplace=False)<br>
)<br>
)<br>
)<br>
(ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)<br>
)<br>
(lm_head): Linear(in_features=768, out_features=50257, bias=False)<br>
)</p>
<p>Also, I print the grad sample shape for each layer, it looks like that a embedding layer has a wrong shape:<br>
params shape: torch.Size([50257, 768]), params grad_sample shape: torch.Size([13, 50257, 768])<br>
params shape: torch.Size([1024, 768]), params grad_sample shape: torch.Size([1, 1024, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 2304]), params grad_sample shape: torch.Size([13, 768, 2304])<br>
params shape: torch.Size([2304]), params grad_sample shape: torch.Size([13, 2304])<br>
params shape: torch.Size([768, 768]), params grad_sample shape: torch.Size([13, 768, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 3072]), params grad_sample shape: torch.Size([13, 768, 3072])<br>
params shape: torch.Size([3072]), params grad_sample shape: torch.Size([13, 3072])<br>
params shape: torch.Size([3072, 768]), params grad_sample shape: torch.Size([13, 3072, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 2304]), params grad_sample shape: torch.Size([13, 768, 2304])<br>
params shape: torch.Size([2304]), params grad_sample shape: torch.Size([13, 2304])<br>
params shape: torch.Size([768, 768]), params grad_sample shape: torch.Size([13, 768, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 3072]), params grad_sample shape: torch.Size([13, 768, 3072])<br>
params shape: torch.Size([3072]), params grad_sample shape: torch.Size([13, 3072])<br>
params shape: torch.Size([3072, 768]), params grad_sample shape: torch.Size([13, 3072, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 2304]), params grad_sample shape: torch.Size([13, 768, 2304])<br>
params shape: torch.Size([2304]), params grad_sample shape: torch.Size([13, 2304])<br>
params shape: torch.Size([768, 768]), params grad_sample shape: torch.Size([13, 768, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 3072]), params grad_sample shape: torch.Size([13, 768, 3072])<br>
params shape: torch.Size([3072]), params grad_sample shape: torch.Size([13, 3072])<br>
params shape: torch.Size([3072, 768]), params grad_sample shape: torch.Size([13, 3072, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 2304]), params grad_sample shape: torch.Size([13, 768, 2304])<br>
params shape: torch.Size([2304]), params grad_sample shape: torch.Size([13, 2304])<br>
params shape: torch.Size([768, 768]), params grad_sample shape: torch.Size([13, 768, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 3072]), params grad_sample shape: torch.Size([13, 768, 3072])<br>
params shape: torch.Size([3072]), params grad_sample shape: torch.Size([13, 3072])<br>
params shape: torch.Size([3072, 768]), params grad_sample shape: torch.Size([13, 3072, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 2304]), params grad_sample shape: torch.Size([13, 768, 2304])<br>
params shape: torch.Size([2304]), params grad_sample shape: torch.Size([13, 2304])<br>
params shape: torch.Size([768, 768]), params grad_sample shape: torch.Size([13, 768, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 3072]), params grad_sample shape: torch.Size([13, 768, 3072])<br>
params shape: torch.Size([3072]), params grad_sample shape: torch.Size([13, 3072])<br>
params shape: torch.Size([3072, 768]), params grad_sample shape: torch.Size([13, 3072, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 2304]), params grad_sample shape: torch.Size([13, 768, 2304])<br>
params shape: torch.Size([2304]), params grad_sample shape: torch.Size([13, 2304])<br>
params shape: torch.Size([768, 768]), params grad_sample shape: torch.Size([13, 768, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768, 3072]), params grad_sample shape: torch.Size([13, 768, 3072])<br>
params shape: torch.Size([3072]), params grad_sample shape: torch.Size([13, 3072])<br>
params shape: torch.Size([3072, 768]), params grad_sample shape: torch.Size([13, 3072, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])<br>
params shape: torch.Size([768]), params grad_sample shape: torch.Size([13, 768])</p>
<p>How can I solve this problem? Any answer is greatly appreciated.</p>