fix smart schedule shadow sync bug

laekov · laekov · commit dfd9b5d5ce3f · 2023-09-11T19:11:03.000+08:00
diff --git a/cuda/fastermoe/smart_schedule.h b/cuda/fastermoe/smart_schedule.h
@@ -338,7 +338,7 @@ void fmoe_cuda_fused_backward_impl(
             collect_fn(si, i / num_expert, 0);
             if (i / num_expert == rank) {
                 cudaEventCreate(evt_reduce + i % num_expert);
-                cudaEventRecord(evt_reduce[i % num_expert], smgr->stream(num_expert));
+                cudaEventRecord(evt_reduce[i % num_expert], smgr->stream(0));
             }
             ++si;
         }
@@ -367,7 +367,6 @@ void fmoe_cuda_fused_backward_impl(
     for (long i = 0, si = 0; i < world_size * num_expert; ++i) {
         if (stored_models[i]) {
             if (i / num_expert == rank) {
-                FMOE_SWE(smgr->stream(0), evt_reduce[i % num_expert]);
                 FMOE_SWE(smgr->torchStream(), evt_reduce[i % num_expert]);
                 set_grad_fn(si, i % num_expert);
             }

Original file line number	Diff line number	Diff line change
`@@ -338,7 +338,7 @@ void fmoe_cuda_fused_backward_impl(`
`338`	`338`	`collect_fn(si, i / num_expert, 0);`
`339`	`339`	`if (i / num_expert == rank) {`
`340`	`340`	`cudaEventCreate(evt_reduce + i % num_expert);`
`341`		`- cudaEventRecord(evt_reduce[i % num_expert], smgr->stream(num_expert));`
	`341`	`+ cudaEventRecord(evt_reduce[i % num_expert], smgr->stream(0));`
`342`	`342`	`}`
`343`	`343`	`++si;`
`344`	`344`	`}`
`@@ -367,7 +367,6 @@ void fmoe_cuda_fused_backward_impl(`
`367`	`367`	`for (long i = 0, si = 0; i < world_size * num_expert; ++i) {`
`368`	`368`	`if (stored_models[i]) {`
`369`	`369`	`if (i / num_expert == rank) {`
`370`		`- FMOE_SWE(smgr->stream(0), evt_reduce[i % num_expert]);`
`371`	`370`	`FMOE_SWE(smgr->torchStream(), evt_reduce[i % num_expert]);`
`372`	`371`	`set_grad_fn(si, i % num_expert);`
`373`	`372`	`}`