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!

[TianBian95]

In fact, using tgt_shape = (-1, bsz // self.device_count, *corr_x.shape[2:]) and corr_x = corr_x.reshape(tgt_shape) has a certain probability of throwing an error. For example, with the default settings of 200 samples and a batch size of 32 per device, if self.device_count=6, all data will be divided into 2 batches. The first bsz will be 32*6=192, and the second bsz will be 200-192=8. At this point, tgt_shape will be (1264, 1, 64, 768), and each device's corr_x will be (211, 1, 64, 768), which will raise an untimeError at line 569 in modeling_fpt2.py: einsum(): subscript w has size 158 for operand 1 which does not broadcast with previously seen size 211. Of course, with 2 or 4 GPUs, this bug will be missed.