RuntimeError: Tensor Size Mismatch During Training

[XiaoXiong-Sherry]

I am encountering the following error while executing run_scripts/gt_sweep.sh:
File "Edge-Pruning/src/modeling/modeling_fpt2.py", line 1293, in forward hidden_states, embeds, z_nodes_sum = self.write(inputs_embeds, position_embeds, corr_x=corr_x)
File "Edge-Pruning/src/modeling/modeling_fpt2.py", line 1184, in write tok_embeds = tok_embeds + corr_x[0] * (1 - z_tokens)
RuntimeError: The size of tensor a (32) must match the size of tensor b (64) at non-singleton dimension 0
Thank you for your assistance!

[testzer0]

Hi Xiao Xiong,
Thank you for bringing this to our attention! I ran the gt_sweep for a sparsity of 0.95 on my end, and did not encounter this error. Could you provide more details, such as which sparsity you got this error for, and whether it popped up at the beginning of training or a while later? If it helps, I am using transformers==4.40.0.dev0, so maybe give that a go? Thanks!

[ZekaiZhaostats]

I also have a similar issue. When I run bash ioi_sweep.sh, it returns me that src/modeling/modeling_fpt2.py", line 1192, in write tok_embeds = tok_embeds + corr_x[0] * (1 - z_tokens)
RuntimeError: The size of tensor a (16) must match the size of tensor b (128) at non-singleton dimension 0.

[TianBian95]

@XiaoXiong-Sherry @ZekaiZhaostats I also have this issue when I run ioi_sweep.sh with two GPU device, the dims of corr_x is (79,64,64,768), the batch size is not 32 since distributed calculation. I set os.environ['CUDA_VISIBLE_DEVICES']='0' to solve the issue. Maybe reshape corr_x as (158, 32, 64, 768) will also solve the issue. @testzer0

[testzer0]

Thank you all for the suggestions! I ran it with two GPUs and indeed, I got the error originally reported. It seems to have gone undetected since we ran our GPT-2 experiments on one GPU. Based on the suggestion by @TianBian95, I have now introduced a fix to the fpt2_* scripts. Hopefully, it works now!