Skip to content

feat(metax): FSDP2 full feature parity + Qwen3 training match #62

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhaoyinglia merged 1 commit into from
Aug 6, 2026
Merged

feat(metax): FSDP2 full feature parity + Qwen3 training match #62

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 52 additions & 0 deletions csrc/aten/copy_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,40 @@
#include "backends/musa/mudnn_common.h"
#endif

// On the CUDA-family backends (including MetaX boxing) the flagos device shares
// the vendor's CUDA streams, so the current stream is readable from c10::cuda.
#if !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_GCU) && \
!defined(USE_MUSA)
#define FLAGOS_COPY_HAS_CUDA_STREAM 1
#include <c10/cuda/CUDAStream.h>
#endif

namespace at::native::flagos {

namespace {

// `Memcpy` is the *synchronous* cudaMemcpy, which only orders against the
// legacy default stream. PyTorch creates its side streams with
// cudaStreamNonBlocking, so work enqueued on one is NOT awaited by a plain
// cudaMemcpy: the copy can read a buffer before the kernel producing it has
// run. That is a silent wrong-data bug, not a crash -- FSDP2 CPU offload hit it
// because its gradient D2H happens inside the reduce-scatter stream, and the
// gradient reached the CPU as zeros (or as a previous tensor's contents).
//
// Synchronizing the current stream before a blocking copy restores the
// semantics callers expect from a synchronous memcpy. It is a no-op on the
// default stream (already ordered), so the common path is unaffected.
inline void SyncCurrentStreamBeforeBlockingCopy() {
#if defined(FLAGOS_COPY_HAS_CUDA_STREAM)
auto stream = c10::cuda::getCurrentCUDAStream();
if (stream.stream() != nullptr) {
stream.synchronize();
}
#endif
}

} // namespace

ADD_IMPL_TO_DISPATCHER(
LocalScalarDenseFn, local_scalar_dense_dispatcher, "_local_scalar_dense")
ADD_IMPL_TO_DISPATCHER(ToCopyFn, to_copy_dispatcher, "_to_copy")
Expand Down Expand Up @@ -73,6 +105,7 @@ at::Tensor _copy_from(
// Fast path: both contiguous, same shape and dtype → direct memcpy.
size_t nbytes = self.numel() * self.element_size();
if (nbytes > 0) {
SyncCurrentStreamBeforeBlockingCopy();
Memcpy(dst.data_ptr(), self.data_ptr(), nbytes, MemcpyDeviceToDevice);
}
} else {
Expand Down Expand Up @@ -153,6 +186,12 @@ at::Tensor _copy_from(

size_t nbytes = self_contig.numel() * self_contig.element_size();

// Every branch below issues a blocking `Memcpy`. The source may have been
// produced by kernels on a non-default stream (and `contiguous()` above may
// itself have just enqueued one there), which a synchronous cudaMemcpy does
// not wait for. Drain the current stream first.
SyncCurrentStreamBeforeBlockingCopy();

if (self.is_cpu() && dst.is_privateuseone()) {
if (dst.is_contiguous()) {
Memcpy(dst.data_ptr(), self_contig.data_ptr(), nbytes, MemcpyHostToDevice);
Expand Down Expand Up @@ -221,6 +260,9 @@ at::Scalar _local_scalar_dense(const at::Tensor& self) {
self.numel() == 1,
"_local_scalar_dense expects a tensor with 1 element");
at::Tensor cpu_tensor = at::empty({1}, self.options().device(at::kCPU));
// `.item()` on a value just computed on a side stream must see that kernel's
// result, and a blocking cudaMemcpy does not wait for a non-blocking stream.
SyncCurrentStreamBeforeBlockingCopy();
Memcpy(
cpu_tensor.data_ptr(),
self.data_ptr(),
Expand Down Expand Up @@ -293,6 +335,12 @@ at::Tensor _to_copy(

at::Tensor result;

// Branches that end in a blocking `Memcpy` drain the current stream first --
// after their `contiguous()` call, which may itself enqueue a kernel there.
// `tensor.to("cpu")` inside `with flagos.stream(s)` returned stale data
// before this. Branches that redispatch to a native CUDA kernel instead need
// no drain: those kernels are stream-ordered on their own.

if (src_is_flagos && dst_is_cuda) {
int device_index =
device.index() >= 0 ? device.index()
Expand All @@ -308,6 +356,7 @@ at::Tensor _to_copy(
self_contig.options().device(c10::Device(c10::kCUDA, device_index)));
size_t nbytes = self_contig.numel() * self_contig.element_size();
if (nbytes > 0) {
SyncCurrentStreamBeforeBlockingCopy();
Memcpy(temp.data_ptr(), self_contig.data_ptr(), nbytes, MemcpyDeviceToDevice);
}
result = (dtype != self.scalar_type()) ? temp.to(dtype) : temp;
Expand Down Expand Up @@ -389,6 +438,7 @@ at::Tensor _to_copy(
self_contig.options().device(c10::Device(c10::kPrivateUse1, device_index)));
size_t nbytes = self_contig.numel() * self_contig.element_size();
if (nbytes > 0) {
SyncCurrentStreamBeforeBlockingCopy();
Memcpy(
result.data_ptr(),
self_contig.data_ptr(),
Expand All @@ -402,6 +452,7 @@ at::Tensor _to_copy(
at::empty(self_contig.sizes(), self_contig.options().device(at::kCPU));
size_t nbytes = self_contig.numel() * self_contig.element_size();
if (nbytes > 0) {
SyncCurrentStreamBeforeBlockingCopy();
Memcpy(temp.data_ptr(), self_contig.data_ptr(), nbytes, MemcpyDeviceToHost);
}
result = (dtype != self.scalar_type()) ? temp.to(dtype) : temp;
Expand All @@ -416,6 +467,7 @@ at::Tensor _to_copy(
src_contig.options().device(c10::Device(c10::kPrivateUse1, device_index)));
size_t nbytes = src_contig.numel() * src_contig.element_size();
if (nbytes > 0) {
SyncCurrentStreamBeforeBlockingCopy();
if (self.is_cpu()) {
Memcpy(result.data_ptr(), src_contig.data_ptr(), nbytes, MemcpyHostToDevice);
} else if (self.is_cuda()) {
Expand Down
Loading
Loading