Training cluster performance boost for bert models part 2

[jstjohn]

See slack thread https://nvidia.slack.com/archives/C074Z808N05/p1730301123648729

The main issue I was hitting was the setting SBATCH --overcommit in my personal slurm script. As a side effect of this exploration I did a few things though which this PR has:

• Identified that a bug in some combination of NeMo2 and Megatron-lm results in some fused kernels getting regularly re-compiled
• Verified that even with this re-compilation bug, we still have the best performance when using these kernels
• Added an option for geneformer to turn on the torch debugger that errors out if a kernel is getting recompiled, this is how I found which kernels were at fault and where/why the recompilations were happening.

Performance summary of different settings

Name	Replicate	Num GPUs	Time per 10 steps	Average Timing
no_recompile	0	1	2.487	2.50875
no_recompile	0	2	2.527
no_recompile	1	1	2.503
no_recompile	1	2	2.518
fused_bias_act	0	1	2.489	2.496666667
fused_bias_act	0	2	2.514
fused_bias_act	1	1	2.487
fused_bias_act	1	2
fused_bias_act_do	0	1	2.459	2.47425
fused_bias_act_do	0	2	2.478
fused_bias_act_do	1	1	2.471
fused_bias_act_do	1	2	2.489
fused_loss	0	1	2.312	2.326
fused_loss	0	2	2.335
fused_loss	1	1	2.323
fused_loss	1	2	2.334
fused_bias_do	0	1	2.467	2.4845
fused_bias_do	0	2	2.499
fused_bias_do	1	1	2.472
fused_bias_do	1	2	2.5
fused_bias_loss	0	1	2.282	2.28775
fused_bias_loss	0	2	2.297
fused_bias_loss	1	1	2.277
fused_bias_loss	1	2	2.295
fused_bias_loss_arange_expand	0	1	2.277	2.29075
fused_bias_loss_arange_expand	0	2	2.298
fused_bias_loss_arange_expand	1	1	2.285
fused_bias_loss_arange_expand	1	2	2.303

[jstjohn]

/build-ci

[jstjohn]

/build-ci

[jstjohn]

/build-ci

[jstjohn]

/build-ci

[jstjohn]

/build-ci

[jstjohn]

/build-ci

[malcolmgreaves]

LGTM -- only requesting one change to not use a known Python anti-pattern. Everything else is good to go though!

[jstjohn]

/build-ci

[jstjohn]

/build-ci