[DTensor] Add test with parallelize_module#2598

t-vi merged 9 commits intomainLightning-AI/lightning-thunder:mainfrom

dtensor-tp-colwise-testLightning-AI/lightning-thunder:dtensor-tp-colwise-testCopy head branch name to clipboard