Allow linear to be consumed by nvFuser by default

[IvanYashchuk]

This change lowers peak memory usage of LitGPT implementations that use mlp_class_name="GptNeoxMLP" configuration (#1175, #1233, #246).

config (mlp_class_name)	Before	This PR
stablecode-completion-alpha-3b (GptNeoxMLP)	889.68 ms | 77.02 GB	892.37 ms | 74.33 GB
Llama-2-7b-hf (LLaMAMLP)	336.06 ms | 64.22 GB	340.86 ms | 64.18 GB

Better memory usage comes from simplifying the setup for Thunder's fusion rematerialization. With this change, there are fewer "producer" fusions.

cc @Priya2698, @wujingyue

[wujingyue]

There are multiple concerning CI failures:

1. FAILED thunder/tests/test_nvfuser.py::test_cse_rematerialization_nvfuser_cuda_None - assert 3 == 11 is likely due to a behavior change in rematerialization that should be reflected as well in the test.
2. FAILED thunder/tests/test_jit_general.py::test_litgpt_variants_kvcache[cuda-llama1-like] - RuntimeError: !detect_exception_in_thread_pool.load() INTERNAL ASSERT FAILED at "/workspace/Fuser/csrc/kernel_cache.cpp":1234, please report a bug with repro script to NVFuser at https://github.com/NVIDIA/Fuser/issues. Detected exception while compiling fusion segments in parallel. Error messages from all threads are printed below. sounds like a bug in nvFuser.

Error from segmentation group 4: producer->getMemoryType() == MemoryType::Global INTERNAL ASSERT FAILED at "/workspace/Fuser/csrc/device_lower/analysis/sync_information.cpp":699, please report a bug with repro script to NVFuser at https://github.com/NVIDIA/Fuser/issues. Inconsistent parallelization found between TV30 (T30_l[ iblockIdx.x284{( ceilDiv(( ceilDiv(( ceilDiv(( 1 * 16 ), 128) ), 1) ), 1) )}, iUS285{1}, iS283{1}, ithreadIdx.x281{128} ]) and TV25(T25_l[ iblockIdx.x232{( ceilDiv(( ceilDiv(( ceilDiv(( 1 * ( 1 * ( 1 * 16 ) ) ), 128) ), 1) ), 1) )}, iUS233{1}, iS231{1}, ithreadIdx.x229{128} ] ca_pos( 4 )). Producer is required to be in Global Memory based on parallelization strategy. RAW flags: (blockIdx.x threadIdx.x)

[Priya2698]

We did not run Thunder benchmarks using nvfuser linear. Should we run other benchmarks as well before enabling it by default?

Additionally, @wujingyue needed to remove support for 1D weights to facilitate DID-aware execution. We might have to add an additional check on Thunder side or use unsqueeze-squeeze operators to support it. I believe we currently do not have cases exercising 1D weights in Thunder so this should not break anything right now.

[crcrpar]

orthogonal to this PR but related to nvfuser knobs: are there enable_foobars other than linear and matmul?

[IvanYashchuk]

are there enable_foobars other than linear and matmul?

Yes, there's also nv_enable_sdpa. I didn't enable matmul in this PR because it leads to excessive memory usage in my tests. I haven't tried nv_enable_sdpa I thin we want to continue using the one from cuDNN for performance.

[wujingyue]

CI with torch-nightly is now passing with NVIDIA/Fuser#3369 fixed. CI with older versions of torch (and therefore older versions of nvFuser) still fail, because we can't fix a past version. @IvanYashchuk, can you up nvFuser's version and enable linear only for that?

[IvanYashchuk]

Yes, I will do that.

[wujingyue]

I'm unsure about the lightning-thunder (ipynb) test. Is that an infra failure or indeed a regression?

[IvanYashchuk]

For some reason, there was a timeout in the cell execution. Let's try rerunning the failed job.

[IvanYashchuk]

@kiya00, could you please check what's going on with the Notebooks CI job?