Add additional carveout for DP

Signed-off-by: rachitg <[email protected]>

transformer_engine/pytorch/csrc/comm_gemm_overlap.h

@@ -86,6 +86,12 @@ struct UbufCommOverlap : torch::CustomClassHolder, UbufBase {
     cga_size = comm_cga_size;
     _empty_tensor = empty_tensor;
 
+    const char* ext_margin_sm = std::getenv("NVTE_EXT_MARGIN_SM");
+
+    int num_ext_margin_sm = 0;
+    if (ext_margin_sm != NULL){
+        num_ext_margin_sm = atoi(ext_margin_sm);
+    }
     // Allocate and register extra userbuffers
     int ubuf_bytes = sample.numel() * sample.element_size();
     _ub_reg = register_user_buffer_collective(reinterpret_cast<void **>(&_ubuf_ptr), ubuf_bytes,
@@ -112,6 +118,7 @@ struct UbufCommOverlap : torch::CustomClassHolder, UbufBase {
     cudaDeviceProp prop;
     cudaGetDeviceProperties(&prop, 0);
     _math_sms = (set_sm_margin) ? prop.multiProcessorCount - num_comm_sm : prop.multiProcessorCount;
+    _math_sms -= num_ext_margin_sm;
 
     output_tensor = torch::Tensor();
     auto counter_options = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA);
@@ -555,6 +562,12 @@ struct UbufP2PCommOverlap : torch::CustomClassHolder, UbufBase {
     sms = 1;
     cga_size = 1;
 
+    const char* ext_margin_sm = std::getenv("NVTE_EXT_MARGIN_SM");
+    int num_ext_margin_sm = 0;
+    if (ext_margin_sm != NULL){
+        num_ext_margin_sm = atoi(ext_margin_sm);
+    }
+
     _empty_tensor = empty_tensor;
     // Create workspace tensor with userbuffer
     int ubuf_bytes = sample.numel() * sample.element_size();
@@ -587,6 +600,7 @@ struct UbufP2PCommOverlap : torch::CustomClassHolder, UbufBase {
     cudaDeviceProp prop;
     cudaGetDeviceProperties(&prop, 0);
     _math_sms = (set_sm_margin) ? prop.multiProcessorCount - num_comm_sm : prop.multiProcessorCount;
+    _math_sms -= num_ext_margin_sm;
 
     _tp_size = tp_size;
     _aggregate2 = aggregate2;