This is I think the same question from here Question about tensor parallel (DTensor, parallelize_module)