diff --git a/csrc/aten/copy_ops.cc b/csrc/aten/copy_ops.cc
index ddb19532..eb56e719 100644
--- a/csrc/aten/copy_ops.cc
+++ b/csrc/aten/copy_ops.cc
@@ -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")
@@ -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 {
@@ -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);
@@ -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(),
@@ -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()
@@ -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;
@@ -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(),
@@ -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;
@@ -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()) {
diff --git a/csrc/aten/generated/cuda_kernels.cc b/csrc/aten/generated/cuda_kernels.cc
index de37acd0..47c31cbe 100644
--- a/csrc/aten/generated/cuda_kernels.cc
+++ b/csrc/aten/generated/cuda_kernels.cc
@@ -948,6 +948,7 @@ void PrivAmpForeachNonFiniteCheckAndUnscaleOutKernelCuda(at::TensorList self, at
   guard.box(self_vec);
   guard.box({found_inf});
   guard.box(out_vec);
+  guard.box({inv_scale});
   at::_amp_foreach_non_finite_check_and_unscale_outf(self_vec, found_inf, inv_scale, out_vec);
 }
 
@@ -955,6 +956,7 @@ void PrivAmpForeachNonFiniteCheckAndUnscaleInplaceKernelCuda(at::TensorList self
   auto self_vec = MaterializeToTensorVec(self);
   TensorListBoxingGuard guard;
   guard.box(self_vec);
+  guard.box({inv_scale});
   at::_amp_foreach_non_finite_check_and_unscale_(self_vec, found_inf, inv_scale);
 }
 
@@ -1378,6 +1380,16 @@ void PrivCudnnRnnBackwardOutKernelCuda(const at::Tensor & input, at::TensorList
   guard.box({out1});
   guard.box({out2});
   guard.box(out3_vec);
+  guard.box({input});
+  guard.box({weight_buf});
+  guard.box({hx});
+  if (cx.has_value()) guard.box({*cx});
+  guard.box({output});
+  if (grad_output.has_value()) guard.box({*grad_output});
+  if (grad_hy.has_value()) guard.box({*grad_hy});
+  if (grad_cy.has_value()) guard.box({*grad_cy});
+  if (dropout_state.has_value()) guard.box({*dropout_state});
+  guard.box({reserve});
   at::_cudnn_rnn_backward_outf(input, weight_vec, weight_stride0, weight_buf, hx, cx, output, grad_output, grad_hy, grad_cy, mode, hidden_size, proj_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state, reserve, output_mask, out0, out1, out2, out3_vec);
 }
 
@@ -1795,6 +1807,7 @@ ::std::vector<at::Tensor> ForeachAddTensorKernelCuda(at::TensorList self, const
   auto self_vec = MaterializeToTensorVec(self);
   TensorListBoxingGuard guard;
   guard.box(self_vec);
+  guard.box({other});
   auto result = at::_foreach_add(self_vec, other, alpha);
   UnboxTensorVecToFlagos(result);
   return result;
@@ -1806,6 +1819,7 @@ void ForeachAddTensorOutKernelCuda(at::TensorList self, const at::Tensor & other
   TensorListBoxingGuard guard;
   guard.box(self_vec);
   guard.box(out_vec);
+  guard.box({other});
   at::_foreach_add_outf(self_vec, other, alpha, out_vec);
 }
 
@@ -1836,6 +1850,7 @@ void ForeachAddInplaceTensorKernelCuda(at::TensorList self, const at::Tensor & o
   auto self_vec = MaterializeToTensorVec(self);
   TensorListBoxingGuard guard;
   guard.box(self_vec);
+  guard.box({other});
   at::_foreach_add_(self_vec, other, alpha);
 }
 
@@ -1899,6 +1914,7 @@ ::std::vector<at::Tensor> ForeachAddcdivTensorKernelCuda(at::TensorList self, at
   guard.box(self_vec);
   guard.box(tensor1_vec);
   guard.box(tensor2_vec);
+  guard.box({scalars});
   auto result = at::_foreach_addcdiv(self_vec, tensor1_vec, tensor2_vec, scalars);
   UnboxTensorVecToFlagos(result);
   return result;
@@ -1914,6 +1930,7 @@ void ForeachAddcdivTensorOutKernelCuda(at::TensorList self, at::TensorList tenso
   guard.box(tensor1_vec);
   guard.box(tensor2_vec);
   guard.box(out_vec);
+  guard.box({scalars});
   at::_foreach_addcdiv_outf(self_vec, tensor1_vec, tensor2_vec, scalars, out_vec);
 }
 
@@ -1947,6 +1964,7 @@ void ForeachAddcdivInplaceTensorKernelCuda(at::TensorList self, at::TensorList t
   guard.box(self_vec);
   guard.box(tensor1_vec);
   guard.box(tensor2_vec);
+  guard.box({scalars});
   at::_foreach_addcdiv_(self_vec, tensor1_vec, tensor2_vec, scalars);
 }
 
@@ -2010,6 +2028,7 @@ ::std::vector<at::Tensor> ForeachAddcmulTensorKernelCuda(at::TensorList self, at
   guard.box(self_vec);
   guard.box(tensor1_vec);
   guard.box(tensor2_vec);
+  guard.box({scalars});
   auto result = at::_foreach_addcmul(self_vec, tensor1_vec, tensor2_vec, scalars);
   UnboxTensorVecToFlagos(result);
   return result;
@@ -2025,6 +2044,7 @@ void ForeachAddcmulTensorOutKernelCuda(at::TensorList self, at::TensorList tenso
   guard.box(tensor1_vec);
   guard.box(tensor2_vec);
   guard.box(out_vec);
+  guard.box({scalars});
   at::_foreach_addcmul_outf(self_vec, tensor1_vec, tensor2_vec, scalars, out_vec);
 }
 
@@ -2058,6 +2078,7 @@ void ForeachAddcmulInplaceTensorKernelCuda(at::TensorList self, at::TensorList t
   guard.box(self_vec);
   guard.box(tensor1_vec);
   guard.box(tensor2_vec);
+  guard.box({scalars});
   at::_foreach_addcmul_(self_vec, tensor1_vec, tensor2_vec, scalars);
 }
 
@@ -2441,6 +2462,7 @@ ::std::vector<at::Tensor> ForeachDivTensorKernelCuda(at::TensorList self, const
   auto self_vec = MaterializeToTensorVec(self);
   TensorListBoxingGuard guard;
   guard.box(self_vec);
+  guard.box({other});
   auto result = at::_foreach_div(self_vec, other);
   UnboxTensorVecToFlagos(result);
   return result;
@@ -2452,6 +2474,7 @@ void ForeachDivTensorOutKernelCuda(at::TensorList self, const at::Tensor & other
   TensorListBoxingGuard guard;
   guard.box(self_vec);
   guard.box(out_vec);
+  guard.box({other});
   at::_foreach_div_outf(self_vec, other, out_vec);
 }
 
@@ -2482,6 +2505,7 @@ void ForeachDivInplaceTensorKernelCuda(at::TensorList self, const at::Tensor & o
   auto self_vec = MaterializeToTensorVec(self);
   TensorListBoxingGuard guard;
   guard.box(self_vec);
+  guard.box({other});
   at::_foreach_div_(self_vec, other);
 }
 
@@ -3101,6 +3125,7 @@ ::std::vector<at::Tensor> ForeachMulTensorKernelCuda(at::TensorList self, const
   auto self_vec = MaterializeToTensorVec(self);
   TensorListBoxingGuard guard;
   guard.box(self_vec);
+  guard.box({other});
   auto result = at::_foreach_mul(self_vec, other);
   UnboxTensorVecToFlagos(result);
   return result;
@@ -3112,6 +3137,7 @@ void ForeachMulTensorOutKernelCuda(at::TensorList self, const at::Tensor & other
   TensorListBoxingGuard guard;
   guard.box(self_vec);
   guard.box(out_vec);
+  guard.box({other});
   at::_foreach_mul_outf(self_vec, other, out_vec);
 }
 
@@ -3142,6 +3168,7 @@ void ForeachMulInplaceTensorKernelCuda(at::TensorList self, const at::Tensor & o
   auto self_vec = MaterializeToTensorVec(self);
   TensorListBoxingGuard guard;
   guard.box(self_vec);
+  guard.box({other});
   at::_foreach_mul_(self_vec, other);
 }
 
@@ -3692,6 +3719,8 @@ void PrivFusedAdagradOutKernelCuda(at::TensorList self, at::TensorList grads, at
   guard.box(state_sums_vec);
   guard.box(state_steps_vec);
   guard.box(out_vec);
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adagrad_outf(self_vec, grads_vec, state_sums_vec, state_steps_vec, lr, lr_decay, weight_decay, eps, maximize, grad_scale, found_inf, out_vec);
 }
 
@@ -3707,6 +3736,9 @@ void PrivFusedAdagradTensorLrOutKernelCuda(at::TensorList self, at::TensorList g
   guard.box(state_sums_vec);
   guard.box(state_steps_vec);
   guard.box(out_vec);
+  guard.box({lr});
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adagrad_outf(self_vec, grads_vec, state_sums_vec, state_steps_vec, lr, lr_decay, weight_decay, eps, maximize, grad_scale, found_inf, out_vec);
 }
 
@@ -3720,6 +3752,8 @@ void PrivFusedAdagradInplaceKernelCuda(at::TensorList self, at::TensorList grads
   guard.box(grads_vec);
   guard.box(state_sums_vec);
   guard.box(state_steps_vec);
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adagrad_(self_vec, grads_vec, state_sums_vec, state_steps_vec, lr, lr_decay, weight_decay, eps, maximize, grad_scale, found_inf);
 }
 
@@ -3733,6 +3767,9 @@ void PrivFusedAdagradInplaceTensorLrKernelCuda(at::TensorList self, at::TensorLi
   guard.box(grads_vec);
   guard.box(state_sums_vec);
   guard.box(state_steps_vec);
+  guard.box({lr});
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adagrad_(self_vec, grads_vec, state_sums_vec, state_steps_vec, lr, lr_decay, weight_decay, eps, maximize, grad_scale, found_inf);
 }
 
@@ -3752,6 +3789,8 @@ void PrivFusedAdamOutKernelCuda(at::TensorList self, at::TensorList grads, at::T
   guard.box(max_exp_avg_sqs_vec);
   guard.box(state_steps_vec);
   guard.box(out_vec);
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adam_outf(self_vec, grads_vec, exp_avgs_vec, exp_avg_sqs_vec, max_exp_avg_sqs_vec, state_steps_vec, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf, out_vec);
 }
 
@@ -3771,6 +3810,9 @@ void PrivFusedAdamTensorLrOutKernelCuda(at::TensorList self, at::TensorList grad
   guard.box(max_exp_avg_sqs_vec);
   guard.box(state_steps_vec);
   guard.box(out_vec);
+  guard.box({lr});
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adam_outf(self_vec, grads_vec, exp_avgs_vec, exp_avg_sqs_vec, max_exp_avg_sqs_vec, state_steps_vec, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf, out_vec);
 }
 
@@ -3788,6 +3830,8 @@ void PrivFusedAdamInplaceKernelCuda(at::TensorList self, at::TensorList grads, a
   guard.box(exp_avg_sqs_vec);
   guard.box(max_exp_avg_sqs_vec);
   guard.box(state_steps_vec);
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adam_(self_vec, grads_vec, exp_avgs_vec, exp_avg_sqs_vec, max_exp_avg_sqs_vec, state_steps_vec, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
 }
 
@@ -3805,6 +3849,9 @@ void PrivFusedAdamInplaceTensorLrKernelCuda(at::TensorList self, at::TensorList
   guard.box(exp_avg_sqs_vec);
   guard.box(max_exp_avg_sqs_vec);
   guard.box(state_steps_vec);
+  guard.box({lr});
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adam_(self_vec, grads_vec, exp_avgs_vec, exp_avg_sqs_vec, max_exp_avg_sqs_vec, state_steps_vec, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
 }
 
@@ -3824,6 +3871,8 @@ void PrivFusedAdamwOutKernelCuda(at::TensorList self, at::TensorList grads, at::
   guard.box(max_exp_avg_sqs_vec);
   guard.box(state_steps_vec);
   guard.box(out_vec);
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adamw_outf(self_vec, grads_vec, exp_avgs_vec, exp_avg_sqs_vec, max_exp_avg_sqs_vec, state_steps_vec, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf, out_vec);
 }
 
@@ -3843,6 +3892,9 @@ void PrivFusedAdamwTensorLrOutKernelCuda(at::TensorList self, at::TensorList gra
   guard.box(max_exp_avg_sqs_vec);
   guard.box(state_steps_vec);
   guard.box(out_vec);
+  guard.box({lr});
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adamw_outf(self_vec, grads_vec, exp_avgs_vec, exp_avg_sqs_vec, max_exp_avg_sqs_vec, state_steps_vec, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf, out_vec);
 }
 
@@ -3860,6 +3912,8 @@ void PrivFusedAdamwInplaceKernelCuda(at::TensorList self, at::TensorList grads,
   guard.box(exp_avg_sqs_vec);
   guard.box(max_exp_avg_sqs_vec);
   guard.box(state_steps_vec);
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adamw_(self_vec, grads_vec, exp_avgs_vec, exp_avg_sqs_vec, max_exp_avg_sqs_vec, state_steps_vec, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
 }
 
@@ -3877,6 +3931,9 @@ void PrivFusedAdamwInplaceTensorLrKernelCuda(at::TensorList self, at::TensorList
   guard.box(exp_avg_sqs_vec);
   guard.box(max_exp_avg_sqs_vec);
   guard.box(state_steps_vec);
+  guard.box({lr});
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_adamw_(self_vec, grads_vec, exp_avgs_vec, exp_avg_sqs_vec, max_exp_avg_sqs_vec, state_steps_vec, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
 }
 
@@ -3958,6 +4015,8 @@ void PrivFusedSgdOutKernelCuda(at::TensorList self, at::TensorList grads, at::Te
   guard.box(grads_vec);
   guard.box(momentum_buffer_list_vec);
   guard.box(out_vec);
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_sgd_outf(self_vec, grads_vec, momentum_buffer_list_vec, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale, found_inf, out_vec);
 }
 
@@ -3971,6 +4030,9 @@ void PrivFusedSgdTensorLrOutKernelCuda(at::TensorList self, at::TensorList grads
   guard.box(grads_vec);
   guard.box(momentum_buffer_list_vec);
   guard.box(out_vec);
+  guard.box({lr});
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_sgd_outf(self_vec, grads_vec, momentum_buffer_list_vec, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale, found_inf, out_vec);
 }
 
@@ -3982,6 +4044,8 @@ void PrivFusedSgdInplaceKernelCuda(at::TensorList self, at::TensorList grads, at
   guard.box(self_vec);
   guard.box(grads_vec);
   guard.box(momentum_buffer_list_vec);
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_sgd_(self_vec, grads_vec, momentum_buffer_list_vec, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale, found_inf);
 }
 
@@ -3993,6 +4057,9 @@ void PrivFusedSgdInplaceTensorLrKernelCuda(at::TensorList self, at::TensorList g
   guard.box(self_vec);
   guard.box(grads_vec);
   guard.box(momentum_buffer_list_vec);
+  guard.box({lr});
+  if (grad_scale.has_value()) guard.box({*grad_scale});
+  if (found_inf.has_value()) guard.box({*found_inf});
   at::_fused_sgd_(self_vec, grads_vec, momentum_buffer_list_vec, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale, found_inf);
 }
 
@@ -4073,6 +4140,7 @@ at::Tensor PrivJaggedToPaddedDenseForwardKernelCuda(const at::Tensor & values, a
   auto offsets_vec = MaterializeToTensorVec(offsets);
   TensorListBoxingGuard guard;
   guard.box(offsets_vec);
+  guard.box({values});
   auto result = at::_jagged_to_padded_dense_forward(values, offsets_vec, max_lengths, padding_value);
   UnboxToFlagos(result);
   return result;
@@ -4614,6 +4682,7 @@ at::Tensor PrivPaddedDenseToJaggedForwardKernelCuda(const at::Tensor & dense, at
   auto offsets_vec = MaterializeToTensorVec(offsets);
   TensorListBoxingGuard guard;
   guard.box(offsets_vec);
+  guard.box({dense});
   auto result = at::_padded_dense_to_jagged_forward(dense, offsets_vec, total_L);
   UnboxToFlagos(result);
   return result;
@@ -4785,6 +4854,10 @@ at::Tensor PrivScaledGroupedMmV2KernelCuda(const at::Tensor & self, const at::Te
   TensorListBoxingGuard guard;
   guard.box(scale_a_vec);
   guard.box(scale_b_vec);
+  guard.box({self});
+  guard.box({mat2});
+  if (offs.has_value()) guard.box({*offs});
+  if (bias.has_value()) guard.box({*bias});
   auto result = at::_scaled_grouped_mm_v2(self, mat2, scale_a_vec, recipe_a, swizzle_a, scale_b_vec, recipe_b, swizzle_b, offs, bias, out_dtype, contraction_dim, use_fast_accum);
   UnboxToFlagos(result);
   return result;
@@ -4814,6 +4887,9 @@ at::Tensor PrivScaledMmV2KernelCuda(const at::Tensor & self, const at::Tensor &
   TensorListBoxingGuard guard;
   guard.box(scale_a_vec);
   guard.box(scale_b_vec);
+  guard.box({self});
+  guard.box({mat2});
+  if (bias.has_value()) guard.box({*bias});
   auto result = at::_scaled_mm_v2(self, mat2, scale_a_vec, recipe_a, swizzle_a, scale_b_vec, recipe_b, swizzle_b, bias, out_dtype, contraction_dim, use_fast_accum);
   UnboxToFlagos(result);
   return result;
@@ -10984,6 +11060,13 @@ void LstmMpsBackwardOutKernelCuda(const ::std::optional<at::Tensor> & grad_y, co
   guard.box({out0});
   guard.box(out1_vec);
   guard.box(out2_vec);
+  if (grad_y.has_value()) guard.box({*grad_y});
+  if (grad_hy.has_value()) guard.box({*grad_hy});
+  if (grad_cy.has_value()) guard.box({*grad_cy});
+  guard.box({z_state});
+  guard.box({cell_state_fwd});
+  guard.box({input});
+  guard.box({layersOutputs});
   at::lstm_mps_backward_outf(grad_y, grad_hy, grad_cy, z_state, cell_state_fwd, input, layersOutputs, hx_vec, params_vec, has_biases, num_layers, dropout, train, bidirectional, batch_first, out0, out1_vec, out2_vec);
 }
 
@@ -11498,6 +11581,16 @@ void MiopenRnnBackwardOutKernelCuda(const at::Tensor & input, at::TensorList wei
   guard.box({out1});
   guard.box({out2});
   guard.box(out3_vec);
+  guard.box({input});
+  guard.box({weight_buf});
+  guard.box({hx});
+  if (cx.has_value()) guard.box({*cx});
+  guard.box({output});
+  if (grad_output.has_value()) guard.box({*grad_output});
+  if (grad_hy.has_value()) guard.box({*grad_hy});
+  if (grad_cy.has_value()) guard.box({*grad_cy});
+  if (dropout_state.has_value()) guard.box({*dropout_state});
+  guard.box({reserve});
   at::miopen_rnn_backward_outf(input, weight_vec, weight_stride0, weight_buf, hx, cx, output, grad_output, grad_hy, grad_cy, mode, hidden_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state, reserve, output_mask, out0, out1, out2, out3_vec);
 }
 
@@ -15523,6 +15616,7 @@ void SplitCopyTensorOutKernelCuda(const at::Tensor & self, int64_t split_size, i
   auto out_vec = MaterializeToTensorVec(out);
   TensorListBoxingGuard guard;
   guard.box(out_vec);
+  guard.box({self});
   at::split_copy_outf(self, split_size, dim, out_vec);
 }
 
@@ -15537,6 +15631,7 @@ void SplitWithSizesCopyOutKernelCuda(const at::Tensor & self, at::IntArrayRef sp
   auto out_vec = MaterializeToTensorVec(out);
   TensorListBoxingGuard guard;
   guard.box(out_vec);
+  guard.box({self});
   at::split_with_sizes_copy_outf(self, split_sizes, dim, out_vec);
 }
 
@@ -16058,6 +16153,7 @@ void UnbindCopyIntOutKernelCuda(const at::Tensor & self, int64_t dim, at::Tensor
   auto out_vec = MaterializeToTensorVec(out);
   TensorListBoxingGuard guard;
   guard.box(out_vec);
+  guard.box({self});
   at::unbind_copy_outf(self, dim, out_vec);
 }
 
@@ -16170,6 +16266,7 @@ void UnsafeSplitTensorOutKernelCuda(const at::Tensor & self, int64_t split_size,
   auto out_vec = MaterializeToTensorVec(out);
   TensorListBoxingGuard guard;
   guard.box(out_vec);
+  guard.box({self});
   at::unsafe_split_outf(self, split_size, dim, out_vec);
 }
 
@@ -16184,6 +16281,7 @@ void UnsafeSplitWithSizesOutKernelCuda(const at::Tensor & self, at::IntArrayRef
   auto out_vec = MaterializeToTensorVec(out);
   TensorListBoxingGuard guard;
   guard.box(out_vec);
+  guard.box({self});
   at::unsafe_split_with_sizes_outf(self, split_sizes, dim, out_vec);
 }
 
diff --git a/scripts/codegen_ops.py b/scripts/codegen_ops.py
index c9ed7446..32d57185 100644
--- a/scripts/codegen_ops.py
+++ b/scripts/codegen_ops.py
@@ -1484,6 +1484,32 @@ def gen_tuple_return(op, fn_type, ret_type, args, func=None):
 }}"""
 
 
+def _scalar_tensor_box_lines(args) -> str:
+    """`guard.box({...})` lines for the plain const Tensor& args of a foreach op.
+
+    TensorListBoxingGuard kernels used to box only the tensor *lists* (and, for
+    out-variants, the mutable outs). Any remaining `const at::Tensor &` argument
+    stayed on flagos, which is not merely a missed optimization: the at:: call
+    dispatches on *all* its tensor arguments, so one unboxed flagos tensor sends
+    it back to PrivateUse1 -- into this same kernel -- and it recurses until the
+    stack is gone. Mutable outs are skipped here; the caller already boxed them.
+    """
+    lines = ""
+    for t, n in args:
+        if "TensorList" in t or "ITensorListRef" in t:
+            continue
+        if "at::Tensor" not in t:
+            continue
+        if "optional" in t:
+            # e.g. the fused optimizers' grad_scale/found_inf.
+            lines += f"  if ({n}.has_value()) guard.box({{*{n}}});\n"
+            continue
+        if "const" not in t:
+            continue  # mutable out: boxed by the caller
+        lines += f"  guard.box({{{n}}});\n"
+    return lines
+
+
 def gen_foreach(op, fn_type, ret_type, args, func=None):
     """cat + _foreach_*: materialize ITensorListRef, box, call API, unbox result."""
     kn = kernel_name(fn_type)
@@ -1525,6 +1551,12 @@ def gen_foreach(op, fn_type, ret_type, args, func=None):
         mat_name = f"{n}_vec"
         box_lines += f"  guard.box({mat_name});\n"
 
+    # Plain (non-list) Tensor inputs must be boxed too. Boxing only the lists
+    # leaves e.g. _foreach_mul.Tensor's `other` on flagos, and at::_foreach_mul
+    # then re-dispatches to PrivateUse1 -- straight back into this kernel, i.e.
+    # unbounded self-recursion ending in SIGSEGV.
+    box_lines += _scalar_tensor_box_lines(args)
+
     if ret_type == "void":
         body = f"  {api}({call_args_str});"
         return f"""{ret_type} {kn}({args_decl(args)}) {{
@@ -1599,6 +1631,11 @@ def gen_foreach_out(op, fn_type, ret_type, args, func=None):
             if n in mutable_tensors:
                 box_lines += f"  guard.box({{{n}}});\n"
             call_arg_names.append(n)
+    # const Tensor& inputs need boxing as much as the lists and the outs do:
+    # split_with_sizes_copy.out left `self` on flagos, so at::split_with_sizes_
+    # copy_outf re-dispatched to PrivateUse1 and recursed into this kernel until
+    # the stack blew (SIGSEGV). This is FSDP2's all-gather copy-out path.
+    box_lines += _scalar_tensor_box_lines(args)
     call_args_str = ", ".join(call_arg_names)
 
     # Return shape follows ret_type: void (no return) or single `Tensor&`
diff --git a/tests/manual/metax/test_fsdp2_features_metax.py b/tests/manual/metax/test_fsdp2_features_metax.py
new file mode 100644
index 00000000..eecb295c
--- /dev/null
+++ b/tests/manual/metax/test_fsdp2_features_metax.py
@@ -0,0 +1,374 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""FSDP2 (``fully_shard``) feature coverage on MetaX, beyond the smoke test.
+
+tests/manual/metax/test_fsdp_live_metax.py establishes that the core path works:
+DeviceMesh builds, parameters become sharded DTensors, and a 3-layer MLP trained
+with plain SGD reproduces the single-GPU loss trajectory. That is necessary but
+narrow. This test covers what a real FSDP2 job additionally depends on:
+
+  * per-layer wrapping -- ``fully_shard`` applied to submodules as well as the
+    root, which is how FSDP2 is actually used (it produces the reshard-after-
+    forward traffic a single root-only call never exercises)
+  * MixedPrecisionPolicy -- bf16 all-gather with fp32 reduce, the standard
+    production config
+  * Adam -- drives the fused/foreach optimizer kernels over DTensor params,
+    a different code path from SGD
+  * clip_grad_norm_ -- needs a cross-mesh norm reduction, and silently produces
+    wrong norms if the partial-to-replicate reduction is broken
+  * state_dict / load_state_dict -- sharded checkpoint round-trip
+  * gradient accumulation with set_requires_gradient_sync(False)
+  * CPU offload
+  * 2D mesh construction (the composability substrate for FSDP+TP)
+
+Every numerical check is against a single-GPU reference computed in the same
+process, not against "it did not crash". Where a feature legitimately changes
+the numbers (bf16), the tolerance is loosened rather than the check dropped.
+
+Run (from repo root):
+    ACCELERATOR=metax FLAGOS_METAX_BOXING=1 \
+    MACA_PATH=/opt/maca METAX_PATH=/opt/maca \
+    LD_LIBRARY_PATH=/opt/maca/lib:/opt/maca/lib64:$LD_LIBRARY_PATH \
+    PYTHONPATH=$PWD \
+    python tests/manual/metax/test_fsdp2_features_metax.py --world-size 4
+"""
+
+import argparse
+import os
+
+# torch_fl MUST be imported before torch: in boxing mode it preloads the maca
+# libtorch_cuda.so and sets GEMS_VENDOR=metax.
+import torch_fl  # noqa: F401
+import torch
+
+if os.environ.get("FLAGOS_DIST_FORCE_NCCL", "0") != "1":
+    try:
+        import flagcx  # noqa: F401  self-registers "flagcx" (metax adaptor)
+    except ImportError:
+        flagcx = None
+else:
+    flagcx = None
+
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.nn as nn
+
+STEPS = 3
+LR = 0.1
+IN_DIM, HIDDEN, OUT_DIM = 16, 32, 4
+BATCH = 8
+
+
+class Net(nn.Module):
+    """Three named blocks, so per-layer fully_shard has something to wrap."""
+
+    def __init__(self):
+        super().__init__()
+        self.block1 = nn.Sequential(nn.Linear(IN_DIM, HIDDEN), nn.ReLU())
+        self.block2 = nn.Sequential(nn.Linear(HIDDEN, HIDDEN), nn.ReLU())
+        self.head = nn.Linear(HIDDEN, OUT_DIM)
+
+    def forward(self, x):
+        return self.head(self.block2(self.block1(x)))
+
+
+def build_model(device):
+    torch.manual_seed(0)
+    return Net().to(device)
+
+
+def step_input(step, device):
+    torch.manual_seed(1000 + step)
+    return torch.randn(BATCH, IN_DIM, device=device)
+
+
+def train(model, device, opt_cls=torch.optim.SGD, steps=STEPS, **opt_kw):
+    opt = opt_cls(model.parameters(), lr=LR, **opt_kw)
+    losses = []
+    for step in range(steps):
+        loss = model(step_input(step, device)).sum()
+        loss.backward()
+        opt.step()
+        opt.zero_grad()
+        losses.append(float(loss.detach()))
+    return losses
+
+
+def close(a, b, tol):
+    return len(a) == len(b) and all(
+        abs(x - y) <= tol * max(1.0, abs(y)) for x, y in zip(a, b)
+    )
+
+
+# --------------------------------------------------------------------------
+# individual feature checks
+# --------------------------------------------------------------------------
+
+
+def check_per_layer_wrap(mesh, dev, ref, results):
+    """fully_shard on each block plus the root -- the real usage pattern."""
+    from torch.distributed.fsdp import fully_shard
+
+    model = build_model(dev)
+    for block in (model.block1, model.block2, model.head):
+        fully_shard(block, mesh=mesh)
+    fully_shard(model, mesh=mesh)
+
+    n_shards = sum(1 for p in model.parameters() if hasattr(p, "to_local"))
+    results.append(("per-layer: all params DTensor", n_shards == 6))
+    losses = train(model, dev)
+    results.append(("per-layer: matches single-GPU", close(losses, ref, 1e-3)))
+    return losses
+
+
+def check_mixed_precision(mesh, dev, ref, results):
+    """bf16 all-gather + fp32 reduce. Numbers shift, so the tolerance widens."""
+    from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
+
+    policy = MixedPrecisionPolicy(
+        param_dtype=torch.bfloat16, reduce_dtype=torch.float32
+    )
+    model = build_model(dev)
+    for block in (model.block1, model.block2, model.head):
+        fully_shard(block, mesh=mesh, mp_policy=policy)
+    fully_shard(model, mesh=mesh, mp_policy=policy)
+
+    losses = train(model, dev)
+    # bf16 has ~3 decimal digits; 5% is loose enough for accumulated drift over
+    # 3 steps but still catches a genuinely broken reduction (which shows up as
+    # order-of-magnitude or sign errors, not 5% drift).
+    results.append(
+        ("mixed-precision bf16: tracks single-GPU", close(losses, ref, 5e-2))
+    )
+    # Compute dtype must actually be bf16, else the policy silently did nothing.
+    params = [p for p in model.parameters() if hasattr(p, "to_local")]
+    results.append(
+        ("mixed-precision: sharded params kept fp32", params[0].dtype == torch.float32)
+    )
+
+
+def check_adam(mesh, dev, results):
+    """Adam over DTensor params -- exercises the foreach optimizer kernels."""
+    from torch.distributed.fsdp import fully_shard
+
+    ref_model = build_model(dev)
+    ref_losses = train(ref_model, dev, opt_cls=torch.optim.Adam)
+
+    model = build_model(dev)
+    for block in (model.block1, model.block2, model.head):
+        fully_shard(block, mesh=mesh)
+    fully_shard(model, mesh=mesh)
+    losses = train(model, dev, opt_cls=torch.optim.Adam)
+    results.append(("Adam: matches single-GPU", close(losses, ref_losses, 1e-3)))
+
+    # foreach=True is the default for Adam on CUDA-like devices; assert it was
+    # not silently downgraded to the for-loop path, which would hide breakage in
+    # the _foreach_* boxing kernels.
+    opt = torch.optim.Adam(model.parameters(), lr=LR)
+    results.append(
+        ("Adam: foreach path active", opt.defaults.get("foreach") is not False)
+    )
+
+
+def check_clip_grad_norm(mesh, dev, results):
+    """Cross-shard grad norm. A broken partial->replicate reduce shows up here."""
+    from torch.distributed.fsdp import fully_shard
+
+    ref_model = build_model(dev)
+    ref_model(step_input(0, dev)).sum().backward()
+    ref_norm = float(torch.nn.utils.clip_grad_norm_(ref_model.parameters(), 1.0))
+
+    model = build_model(dev)
+    for block in (model.block1, model.block2, model.head):
+        fully_shard(block, mesh=mesh)
+    fully_shard(model, mesh=mesh)
+    model(step_input(0, dev)).sum().backward()
+    # FSDP2 has no module-level clip_grad_norm_ (that was FSDP1); the standard
+    # utility is DTensor-aware and returns the norm as a replicated DTensor.
+    norm = float(torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0))
+    results.append(
+        (
+            "clip_grad_norm_: matches single-GPU",
+            abs(norm - ref_norm) <= 1e-3 * max(1.0, ref_norm),
+        )
+    )
+
+
+def check_state_dict(mesh, dev, results):
+    """Sharded state_dict round-trip must restore the exact trajectory."""
+    from torch.distributed.fsdp import fully_shard
+
+    model = build_model(dev)
+    for block in (model.block1, model.block2, model.head):
+        fully_shard(block, mesh=mesh)
+    fully_shard(model, mesh=mesh)
+
+    train(model, dev, steps=1)
+    sd = {k: v.detach().clone() for k, v in model.state_dict().items()}
+    all_dtensor = all(hasattr(v, "to_local") for v in sd.values())
+    results.append(("state_dict: entries are DTensor", all_dtensor))
+
+    # Continue from the saved point, then reload and redo -- same losses.
+    after_save = train(model, dev, steps=2)
+    model.load_state_dict(sd)
+    after_load = train(model, dev, steps=2)
+    results.append(
+        (
+            "state_dict: reload reproduces trajectory",
+            close(after_load, after_save, 1e-4),
+        )
+    )
+
+
+def make_sharded(dev, mesh):
+    from torch.distributed.fsdp import fully_shard
+
+    model = build_model(dev)
+    for block in (model.block1, model.block2, model.head):
+        fully_shard(block, mesh=mesh)
+    fully_shard(model, mesh=mesh)
+    return model
+
+
+def grads_of(model):
+    return [p.grad.to_local().clone() for p in model.parameters() if p.grad is not None]
+
+
+def check_grad_accum(mesh, dev, results):
+    """set_requires_gradient_sync(False) defers the reduce, it must not drop it.
+
+    During the no-sync window FSDP2 keeps the unsharded gradient internally and
+    ``p.grad`` stays None; the reduce happens on the next backward with sync
+    re-enabled, which then holds the sum over every microbatch. Reduce-scatter is
+    linear, so that must equal reducing each microbatch separately and letting
+    ``p.grad`` accumulate -- which is the reference below.
+    """
+    n_micro = 3
+
+    ref = make_sharded(dev, mesh)
+    for step in range(n_micro):
+        ref(step_input(step, dev)).sum().backward()
+    ref_grads = grads_of(ref)
+
+    model = make_sharded(dev, mesh)
+    model.set_requires_gradient_sync(False)
+    for step in range(n_micro - 1):
+        model(step_input(step, dev)).sum().backward()
+    # Documented FSDP2 behaviour: the sharded .grad is not populated until sync.
+    deferred = all(p.grad is None for p in model.parameters())
+    results.append(("grad accumulation: reduce deferred while no-sync", deferred))
+
+    model.set_requires_gradient_sync(True)
+    model(step_input(n_micro - 1, dev)).sum().backward()
+    accum = grads_of(model)
+
+    ok = len(accum) == len(ref_grads) and bool(accum)
+    ok = ok and all(torch.isfinite(g).all().item() for g in accum)
+    ok = ok and all(
+        torch.allclose(a, r, rtol=1e-4, atol=1e-5) for a, r in zip(accum, ref_grads)
+    )
+    results.append(("grad accumulation: sum matches per-step reduce", ok))
+
+
+def check_cpu_offload(mesh, dev, ref, results):
+    from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard
+
+    model = build_model(dev)
+    policy = CPUOffloadPolicy()
+    for block in (model.block1, model.block2, model.head):
+        fully_shard(block, mesh=mesh, offload_policy=policy)
+    fully_shard(model, mesh=mesh, offload_policy=policy)
+    losses = train(model, dev)
+    results.append(("cpu offload: matches single-GPU", close(losses, ref, 1e-3)))
+
+
+def check_2d_mesh(world_size, results):
+    """2D mesh is the substrate for FSDP+TP composition."""
+    from torch.distributed.device_mesh import init_device_mesh
+
+    if world_size < 4:
+        return
+    mesh2d = init_device_mesh(
+        "flagos", (world_size // 2, 2), mesh_dim_names=("dp", "tp")
+    )
+    ok = mesh2d.size() == world_size and mesh2d["dp"].size() == world_size // 2
+    results.append(("2D mesh (dp,tp) built", ok))
+
+
+# --------------------------------------------------------------------------
+
+
+def worker(rank: int, world_size: int):
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29713")
+
+    dev = torch.device(f"flagos:{rank}")
+    torch.cuda.set_device(rank)
+    dist.init_process_group(backend="flagos", rank=rank, world_size=world_size)
+
+    from torch.distributed.device_mesh import init_device_mesh
+
+    mesh = init_device_mesh("flagos", (world_size,))
+    results = []
+
+    if rank == 0:
+        print(
+            f"[setup] flagcx={'yes' if flagcx else 'no'} world_size={world_size}",
+            flush=True,
+        )
+
+    ref = train(build_model(dev), dev)
+    if rank == 0:
+        print(f"[ref] sgd losses={[f'{v:.6f}' for v in ref]}", flush=True)
+
+    checks = (
+        ("per-layer wrap", lambda: check_per_layer_wrap(mesh, dev, ref, results)),
+        ("mixed precision", lambda: check_mixed_precision(mesh, dev, ref, results)),
+        ("adam", lambda: check_adam(mesh, dev, results)),
+        ("clip_grad_norm", lambda: check_clip_grad_norm(mesh, dev, results)),
+        ("state_dict", lambda: check_state_dict(mesh, dev, results)),
+        ("grad accum", lambda: check_grad_accum(mesh, dev, results)),
+        ("cpu offload", lambda: check_cpu_offload(mesh, dev, ref, results)),
+        ("2d mesh", lambda: check_2d_mesh(world_size, results)),
+    )
+    for name, fn in checks:
+        try:
+            fn()
+        except Exception as e:  # noqa: BLE001
+            results.append((f"{name} ran", False))
+            if rank == 0:
+                print(f"[{name}] raised {type(e).__name__}: {e}", flush=True)
+
+    dist.barrier()
+    if rank == 0:
+        for name, ok in results:
+            print(f"[{'OK' if ok else 'FAIL'}] {name}", flush=True)
+        n_fail = sum(1 for _, ok in results if not ok)
+        status = "ALL PASS" if n_fail == 0 else f"{n_fail} FAILED"
+        print(
+            f"=== metax fsdp2 features: {status} ({len(results)} checks) ===",
+            flush=True,
+        )
+    dist.destroy_process_group()
+    if any(not ok for _, ok in results):
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--world-size", type=int, default=4)
+    args = ap.parse_args()
+    mp.set_start_method("spawn", force=True)
+    mp.spawn(worker, args=(args.world_size,), nprocs=args.world_size, join=True)
diff --git a/tests/manual/metax/test_fsdp_live_metax.py b/tests/manual/metax/test_fsdp_live_metax.py
new file mode 100644
index 00000000..f9743fcd
--- /dev/null
+++ b/tests/manual/metax/test_fsdp_live_metax.py
@@ -0,0 +1,248 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Live multi-GPU FSDP test on MetaX (run on the 8xC550 box).
+
+Covers both generations of FSDP on the flagos device:
+
+  * FSDP1 -- ``FullyShardedDataParallel``, FlatParameter-based. Needs the comm
+    collectives plus ``resizePrivateUse1Bytes`` (csrc/runtime/hooks.h) for
+    storage resize during (re)sharding.
+  * FSDP2 -- ``fully_shard``, DTensor-based. Additionally needs a DeviceMesh,
+    which is where two real bugs surfaced:
+
+    1. ``ProcessGroupFlagOS`` never registered its inner backend, so
+       ``pg.group_name`` raised "ProcessGroup name not set" and DeviceMesh could
+       not be constructed at all. Fixed by ``_register_inner_backend``.
+    2. ``split_with_sizes_copy.out`` -- FSDP2's all-gather copy-out -- left its
+       ``self`` input on flagos, so the boxing kernel re-dispatched to
+       PrivateUse1 into itself and recursed until SIGSEGV. Fixed in
+       scripts/codegen_ops.py by boxing const Tensor inputs of the
+       TensorListBoxingGuard kernels.
+
+Correctness is checked against a *single-GPU* reference rather than just
+"it ran": sharding is only right if the sharded model takes the same
+optimization trajectory as the unsharded one. Each rank feeds the model the
+same input, so per-step losses must match the reference to tolerance and must
+agree across ranks.
+
+Run (from repo root):
+    ACCELERATOR=metax FLAGOS_METAX_BOXING=1 \
+    MACA_PATH=/opt/maca METAX_PATH=/opt/maca \
+    LD_LIBRARY_PATH=/opt/maca/lib:/opt/maca/lib64:$LD_LIBRARY_PATH \
+    PYTHONPATH=$PWD \
+    python tests/manual/metax/test_fsdp_live_metax.py --world-size 2
+
+Force the mccl(NCCL) inner backend (skip FlagCX even if installed):
+    ... FLAGOS_DIST_FORCE_NCCL=1 python .../test_fsdp_live_metax.py
+"""
+
+import argparse
+import os
+
+# torch_fl MUST be imported before torch: in boxing mode it preloads the maca
+# libtorch_cuda.so and sets GEMS_VENDOR=metax (which drives the vendor profile).
+import torch_fl  # noqa: F401
+import torch
+
+if os.environ.get("FLAGOS_DIST_FORCE_NCCL", "0") != "1":
+    try:
+        import flagcx  # noqa: F401  self-registers "flagcx" (metax adaptor)
+    except ImportError:
+        flagcx = None
+else:
+    flagcx = None
+
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.nn as nn
+
+STEPS = 3
+LR = 0.1
+IN_DIM, HIDDEN, OUT_DIM = 8, 16, 1
+BATCH = 4
+
+
+def build_model(device):
+    """Identical init on every rank (and for the reference) via a fixed seed."""
+    torch.manual_seed(0)
+    return nn.Sequential(
+        nn.Linear(IN_DIM, HIDDEN),
+        nn.ReLU(),
+        nn.Linear(HIDDEN, OUT_DIM),
+    ).to(device)
+
+
+def step_input(step, device):
+    """Same input on every rank, so losses are directly comparable."""
+    torch.manual_seed(1000 + step)
+    return torch.randn(BATCH, IN_DIM, device=device)
+
+
+def reference_losses(device):
+    """Single-GPU (unsharded) trajectory the sharded runs must reproduce."""
+    model = build_model(device)
+    opt = torch.optim.SGD(model.parameters(), lr=LR)
+    losses = []
+    for step in range(STEPS):
+        loss = model(step_input(step, device)).sum()
+        loss.backward()
+        opt.step()
+        opt.zero_grad()
+        losses.append(float(loss.detach()))
+    return losses
+
+
+def train_losses(model, device):
+    opt = torch.optim.SGD(model.parameters(), lr=LR)
+    losses = []
+    for step in range(STEPS):
+        loss = model(step_input(step, device)).sum()
+        loss.backward()
+        opt.step()
+        opt.zero_grad()
+        losses.append(float(loss.detach()))
+    return losses
+
+
+def _agree_across_ranks(values, device, world_size):
+    """True when every rank produced the same list of losses."""
+    t = torch.tensor(values, device=device)
+    gathered = [torch.zeros(len(values), device=device) for _ in range(world_size)]
+    dist.all_gather(gathered, t)
+    first = gathered[0].cpu()
+    return all(torch.allclose(g.cpu(), first, atol=1e-4) for g in gathered)
+
+
+def run_fsdp1(rank, world_size, dev, ref, results):
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+    model = FSDP(build_model(dev), device_id=dev)
+    losses = train_losses(model, dev)
+    matches = all(abs(a - b) < 1e-3 * max(1.0, abs(b)) for a, b in zip(losses, ref))
+    agree = _agree_across_ranks(losses, dev, world_size)
+    results.append(("fsdp1 matches single-GPU", matches))
+    results.append(("fsdp1 ranks agree", agree))
+    if rank == 0:
+        print(f"[fsdp1] losses={[f'{v:.6f}' for v in losses]}", flush=True)
+        print(f"[fsdp1] ref   ={[f'{v:.6f}' for v in ref]}", flush=True)
+
+
+def run_fsdp2(rank, world_size, dev, ref, results):
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.fsdp import fully_shard
+
+    # Needs pg.group_name -> exercises _register_inner_backend.
+    mesh = init_device_mesh("flagos", (world_size,))
+    results.append(("fsdp2 device_mesh built", mesh.size() == world_size))
+
+    model = build_model(dev)
+    full_numel = sum(p.numel() for p in model.parameters())
+    fully_shard(model, mesh=mesh)
+
+    params = list(model.parameters())
+    is_dtensor = all(hasattr(p, "to_local") for p in params)
+    local_numel = sum(p.to_local().numel() for p in params if hasattr(p, "to_local"))
+    results.append(("fsdp2 params are DTensor", is_dtensor))
+    # Sharded across `world_size` ranks, so each rank holds strictly less than
+    # the whole model (padding keeps it from being exactly full/world_size).
+    results.append(("fsdp2 params sharded", 0 < local_numel < full_numel))
+
+    losses = train_losses(model, dev)
+    matches = all(abs(a - b) < 1e-3 * max(1.0, abs(b)) for a, b in zip(losses, ref))
+    agree = _agree_across_ranks(losses, dev, world_size)
+    results.append(("fsdp2 matches single-GPU", matches))
+    results.append(("fsdp2 ranks agree", agree))
+    if rank == 0:
+        print(
+            f"[fsdp2] full={full_numel} local={local_numel} "
+            f"ptype={type(params[0]).__name__}",
+            flush=True,
+        )
+        print(f"[fsdp2] losses={[f'{v:.6f}' for v in losses]}", flush=True)
+        print(f"[fsdp2] ref   ={[f'{v:.6f}' for v in ref]}", flush=True)
+
+
+def run_split_with_sizes_copy(dev, results):
+    """FSDP2's all-gather copy-out primitive; used to recurse into SIGSEGV."""
+    src = torch.arange(10, dtype=torch.float32, device=dev)
+    out = [torch.empty(4, device=dev), torch.empty(6, device=dev)]
+    torch.split_with_sizes_copy(src, [4, 6], dim=0, out=out)
+    ok = out[0].cpu().tolist() == [0, 1, 2, 3] and out[1].cpu().tolist() == [
+        4,
+        5,
+        6,
+        7,
+        8,
+        9,
+    ]
+    results.append(("split_with_sizes_copy.out", ok))
+
+
+def worker(rank: int, world_size: int):
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29691")
+
+    dev = torch.device(f"flagos:{rank}")
+    torch.cuda.set_device(rank)  # flagos:i shares physical MetaX GPU i
+
+    dist.init_process_group(backend="flagos", rank=rank, world_size=world_size)
+    results = []
+
+    if rank == 0:
+        print(
+            f"[setup] GEMS_VENDOR={os.environ.get('GEMS_VENDOR', '?')} "
+            f"flagcx={'yes' if flagcx else 'no'} world_size={world_size}",
+            flush=True,
+        )
+
+    # group_name is what DeviceMesh needs; assert it directly so a regression
+    # here is reported as itself rather than as a confusing DeviceMesh failure.
+    pg = dist.distributed_c10d._get_default_group()
+    try:
+        has_name = bool(pg.group_name)
+    except RuntimeError:
+        has_name = False
+    results.append(("pg.group_name set", has_name))
+
+    run_split_with_sizes_copy(dev, results)
+
+    ref = reference_losses(dev)
+
+    for name, fn in (("fsdp1", run_fsdp1), ("fsdp2", run_fsdp2)):
+        try:
+            fn(rank, world_size, dev, ref, results)
+        except Exception as e:  # noqa: BLE001
+            results.append((f"{name} ran", False))
+            if rank == 0:
+                print(f"[{name}] raised {type(e).__name__}: {e}", flush=True)
+
+    dist.barrier()
+    if rank == 0:
+        for name, ok in results:
+            print(f"[{'OK' if ok else 'FAIL'}] {name}", flush=True)
+        n_fail = sum(1 for _, ok in results if not ok)
+        status = "ALL PASS" if n_fail == 0 else f"{n_fail} FAILED"
+        print(f"=== metax fsdp live: {status} ({len(results)} checks) ===", flush=True)
+    dist.destroy_process_group()
+    if any(not ok for _, ok in results):
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--world-size", type=int, default=2)
+    args = ap.parse_args()
+    mp.set_start_method("spawn", force=True)
+    mp.spawn(worker, args=(args.world_size,), nprocs=args.world_size, join=True)
diff --git a/tests/manual/metax/test_qwen3_fsdp2_metax.py b/tests/manual/metax/test_qwen3_fsdp2_metax.py
new file mode 100644
index 00000000..85e8197d
--- /dev/null
+++ b/tests/manual/metax/test_qwen3_fsdp2_metax.py
@@ -0,0 +1,321 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Qwen3 FSDP2 training on MetaX: flagos vs the vendor's own torch.
+
+The FSDP2 feature tests use a 3-layer MLP, which says nothing about whether a
+real transformer converges. This runs the same Qwen3 training job twice -- once
+through the flagos device, once through MetaX's native torch on `cuda` -- and
+compares the loss trajectories step by step.
+
+The two runs are made comparable rather than merely similar:
+
+  * identical initial weights (loaded from the same checkpoint, no random init)
+  * identical batches: the dummy text dataset is tokenized on CPU and sliced
+    deterministically per rank, so rank r sees the same tokens in both modes
+  * identical optimizer, lr, step count and shard layout (per-decoder-layer
+    fully_shard plus the root)
+
+Each mode writes its losses to JSON; `--compare` then diffs them. Because both
+runs execute the same fp32 kernels in the same order, the trajectories should
+agree to well within fp32 accumulation noise -- a real divergence shows up as a
+growing gap, not as jitter in the last digits.
+
+Run (from repo root), sequentially so the two runs do not share GPUs:
+
+    export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
+    export MACA_PATH=/opt/maca METAX_PATH=/opt/maca
+    export LD_LIBRARY_PATH=/opt/maca/lib:/opt/maca/lib64:$LD_LIBRARY_PATH
+    export PYTHONPATH=$PWD
+
+    ACCELERATOR=metax FLAGOS_METAX_BOXING=1 \
+        python tests/manual/metax/test_qwen3_fsdp2_metax.py --mode flagos
+    python tests/manual/metax/test_qwen3_fsdp2_metax.py --mode native
+    python tests/manual/metax/test_qwen3_fsdp2_metax.py --compare
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+
+DEFAULT_MODEL = (
+    "/root/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/"
+    "c1899de289a04d12100db370d81485cdf75e47ca"
+)
+OUT_DIR = "/tmp/qwen3_fsdp2_metax"
+
+
+def _parse_args(argv=None):
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--mode", choices=("flagos", "native"))
+    ap.add_argument("--compare", action="store_true")
+    ap.add_argument("--world-size", type=int, default=4)
+    ap.add_argument("--model", default=DEFAULT_MODEL)
+    ap.add_argument("--steps", type=int, default=20)
+    ap.add_argument("--batch-size", type=int, default=2)
+    ap.add_argument("--seq-len", type=int, default=128)
+    ap.add_argument("--lr", type=float, default=1e-4)
+    # fp32 accumulation over 20 steps of a 0.6B model drifts in the last few
+    # digits; a broken shard/reduce shows up far above this.
+    ap.add_argument("--tol", type=float, default=2e-2)
+    return ap.parse_args(argv)
+
+
+_ARGS = _parse_args()
+
+# torch_fl MUST be imported before torch in flagos mode: it preloads the maca
+# libtorch_cuda.so. In native mode it must not be imported at all, so that the
+# run is a true vendor-torch baseline.
+if _ARGS.mode == "flagos":
+    import torch_fl  # noqa: F401
+
+    try:
+        import flagcx  # noqa: F401  self-registers the "flagcx" metax adaptor
+    except ImportError:
+        flagcx = None
+else:
+    flagcx = None
+
+import torch  # noqa: E402
+import torch.distributed as dist  # noqa: E402
+import torch.multiprocessing as mp  # noqa: E402
+
+sys.path.insert(
+    0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "common")
+)
+
+
+def out_path(mode):
+    return os.path.join(OUT_DIR, f"losses_{mode}.json")
+
+
+def build_batches(tokenizer, args, world_size):
+    """Tokenize on CPU once, then hand each rank a fixed, disjoint slice.
+
+    Doing this outside the device code keeps the token ids bit-identical between
+    the two modes -- no device RNG, no dataloader shuffling.
+    """
+    from dummy_dataset import DummyTextDataset
+
+    n_needed = args.steps * args.batch_size * world_size
+    ds = DummyTextDataset(
+        tokenizer, num_samples=max(100, n_needed), max_length=args.seq_len
+    )
+    ids = torch.stack([ds[i]["input_ids"] for i in range(n_needed)])
+    mask = torch.stack([ds[i]["attention_mask"] for i in range(n_needed)])
+    # [steps, world_size, batch, seq]
+    shape = (args.steps, world_size, args.batch_size, args.seq_len)
+    return ids.reshape(shape), mask.reshape(shape)
+
+
+def worker(rank, world_size, args):
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29731")
+
+    flagos = args.mode == "flagos"
+    # flagos:i shares physical GPU i, so the vendor device index is set the same
+    # way in both modes.
+    torch.cuda.set_device(rank)
+    dev = torch.device(f"flagos:{rank}" if flagos else f"cuda:{rank}")
+    dist.init_process_group(
+        backend="flagos" if flagos else "nccl", rank=rank, world_size=world_size
+    )
+    mesh_device = "flagos" if flagos else "cuda"
+
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.fsdp import fully_shard
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    mesh = init_device_mesh(mesh_device, (world_size,))
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    all_ids, all_mask = build_batches(tokenizer, args, world_size)
+
+    torch.manual_seed(0)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model, dtype=torch.float32, attn_implementation="eager"
+    ).to(dev)
+    model.train()
+    model.config.use_cache = False
+
+    # Per-layer wrapping plus the root: the layout a real FSDP2 job uses, and the
+    # one that actually produces reshard-after-forward traffic.
+    for layer in model.model.layers:
+        fully_shard(layer, mesh=mesh)
+    fully_shard(model, mesh=mesh)
+
+    n_dtensor = sum(1 for p in model.parameters() if hasattr(p, "to_local"))
+    n_param = sum(1 for _ in model.parameters())
+    local_numel = sum(
+        p.to_local().numel() if hasattr(p, "to_local") else p.numel()
+        for p in model.parameters()
+    )
+    if rank == 0:
+        print(
+            f"[setup] mode={args.mode} world_size={world_size} "
+            f"flagcx={'yes' if flagcx else 'no'} "
+            f"params={n_param} dtensor={n_dtensor} local_numel={local_numel / 1e6:.2f}M",
+            flush=True,
+        )
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
+
+    local_losses, global_losses, step_times = [], [], []
+    for step in range(args.steps):
+        input_ids = all_ids[step, rank].to(dev)
+        attention_mask = all_mask[step, rank].to(dev)
+        # Mask padding out of the loss. The dataset pads short sentences to
+        # seq_len, so leaving pad ids in the labels makes the curve a measure of
+        # "learned to emit <pad>" (a 13 -> 0.4 cliff in one step) rather than of
+        # language-modelling convergence.
+        labels = input_ids.masked_fill(attention_mask == 0, -100)
+
+        torch.cuda.synchronize()
+        t0 = time.time()
+        out = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            use_cache=False,
+        )
+        loss = out.loss
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad(set_to_none=True)
+        torch.cuda.synchronize()
+        step_times.append(time.time() - t0)
+
+        # The per-rank loss only covers that rank's batch; the mean over ranks is
+        # what a training log reports, and it is what has to match.
+        local = float(loss.detach())
+        buf = torch.tensor([local], device=dev, dtype=torch.float64)
+        dist.all_reduce(buf)
+        glob = float(buf.item()) / world_size
+        local_losses.append(local)
+        global_losses.append(glob)
+        if rank == 0:
+            print(
+                f"[step {step:2d}] loss={glob:.6f} (rank0 {local:.6f}) "
+                f"time={step_times[-1]:.2f}s",
+                flush=True,
+            )
+
+    dist.barrier()
+    if rank == 0:
+        tokens = args.batch_size * args.seq_len * world_size * args.steps
+        os.makedirs(OUT_DIR, exist_ok=True)
+        with open(out_path(args.mode), "w") as f:
+            json.dump(
+                {
+                    "mode": args.mode,
+                    "world_size": world_size,
+                    "steps": args.steps,
+                    "lr": args.lr,
+                    "batch_size": args.batch_size,
+                    "seq_len": args.seq_len,
+                    "local_numel": local_numel,
+                    "global_losses": global_losses,
+                    "rank0_losses": local_losses,
+                    "avg_step_s": sum(step_times) / len(step_times),
+                    "throughput_tok_s": tokens / sum(step_times),
+                },
+                f,
+                indent=2,
+            )
+        drop = global_losses[0] - global_losses[-1]
+        print(
+            f"=== qwen3 fsdp2 {args.mode}: {args.steps} steps, "
+            f"loss {global_losses[0]:.4f} -> {global_losses[-1]:.4f} "
+            f"(drop {drop:.4f}), {tokens / sum(step_times):.1f} tok/s ===",
+            flush=True,
+        )
+    dist.destroy_process_group()
+
+
+def compare(args):
+    runs = {}
+    for mode in ("flagos", "native"):
+        path = out_path(mode)
+        if not os.path.exists(path):
+            print(f"[FAIL] missing {path} -- run --mode {mode} first")
+            return 1
+        with open(path) as f:
+            runs[mode] = json.load(f)
+
+    a, b = runs["flagos"], runs["native"]
+    results = []
+    for key in ("world_size", "steps", "lr", "batch_size", "seq_len", "local_numel"):
+        results.append((f"same {key}", a[key] == b[key]))
+
+    la, lb = a["global_losses"], b["global_losses"]
+    results.append(("same step count", len(la) == len(lb)))
+    n = min(len(la), len(lb))
+
+    print(f"\n{'step':>4} {'flagos':>12} {'native':>12} {'abs diff':>10} {'rel':>9}")
+    worst, worst_step = 0.0, -1
+    for i in range(n):
+        d = abs(la[i] - lb[i])
+        rel = d / max(1e-9, abs(lb[i]))
+        if rel > worst:
+            worst, worst_step = rel, i
+        print(f"{i:>4} {la[i]:>12.6f} {lb[i]:>12.6f} {d:>10.6f} {rel:>9.2e}")
+
+    results.append(
+        (
+            f"loss trajectory within {args.tol:.0e} (worst {worst:.2e} @ step {worst_step})",
+            worst <= args.tol,
+        )
+    )
+    # Convergence, not just agreement: an untrained 0.6B on repeated text must
+    # actually come down, or "matching" would only mean both runs are broken.
+    for mode, run in (("flagos", a), ("native", b)):
+        drop = run["global_losses"][0] - run["global_losses"][-1]
+        results.append((f"{mode}: loss decreased (drop {drop:.4f})", drop > 0.0))
+        results.append(
+            (
+                f"{mode}: all losses finite",
+                all(x == x and abs(x) != float("inf") for x in run["global_losses"]),
+            )
+        )
+
+    print()
+    for name, ok in results:
+        print(f"[{'OK' if ok else 'FAIL'}] {name}")
+    print(
+        f"\nthroughput: flagos {a['throughput_tok_s']:.1f} tok/s, "
+        f"native {b['throughput_tok_s']:.1f} tok/s "
+        f"({a['throughput_tok_s'] / b['throughput_tok_s']:.2f}x)"
+    )
+    n_fail = sum(1 for _, ok in results if not ok)
+    status = "ALL PASS" if n_fail == 0 else f"{n_fail} FAILED"
+    print(f"=== qwen3 fsdp2 flagos vs native: {status} ({len(results)} checks) ===")
+    return 1 if n_fail else 0
+
+
+if __name__ == "__main__":
+    if _ARGS.compare:
+        raise SystemExit(compare(_ARGS))
+    if not _ARGS.mode:
+        raise SystemExit("need --mode {flagos,native} or --compare")
+    mp.set_start_method("spawn", force=True)
+    mp.spawn(
+        worker,
+        args=(_ARGS.world_size, _ARGS),
+        nprocs=_ARGS.world_size,
+        join=True,
+    )
diff --git a/torch_fl/comm/process_group.py b/torch_fl/comm/process_group.py
index f2f0307d..b1c102cd 100644
--- a/torch_fl/comm/process_group.py
+++ b/torch_fl/comm/process_group.py
@@ -172,6 +172,36 @@ def __init__(self, store, rank: int, world_size: int, timeout=None):
         self._store = store
         self._timeout = timeout
         self._view_fn = self._build_inner(store, rank, world_size, timeout)
+        self._register_inner_backend()
+
+    def _register_inner_backend(self) -> None:
+        """Expose the inner backend to ProcessGroup's device->backend map.
+
+        Overriding the collective virtuals is enough for plain c10d calls, but
+        not for anything that reaches for the *group identity*.
+        ``ProcessGroup::setGroupName`` forwards to the registered backends, so
+        with none registered the name is never stored and ``pg.group_name``
+        raises "ProcessGroup name not set". DeviceMesh reads exactly that
+        property, which put DTensor -- and therefore FSDP2 (``fully_shard``) and
+        anything else mesh-based -- out of reach.
+
+        Registering under privateuseone with BackendType.CUSTOM gives the name
+        somewhere to live. It does not divert the collectives: the dispatcher
+        keeps calling this class's overrides, so flagos tensors still go through
+        the vendor view conversion before touching the inner backend.
+
+        Best-effort: an inner backend that is not a c10d ``Backend`` (some
+        FlagCX builds return a bare ProcessGroup) simply leaves the pre-existing
+        behaviour in place rather than failing process-group construction.
+        """
+        inner = getattr(self, "_inner", None)
+        if inner is None:
+            return
+        try:
+            device = torch.device("privateuseone", max(torch.cuda.current_device(), 0))
+            self._register_backend(device, dist.ProcessGroup.BackendType.CUSTOM, inner)
+        except Exception:  # noqa: BLE001  - never block PG creation on this
+            pass
 
     def _build_inner(self, store, rank, world_size, timeout):
         """Create the inner backend and return the view-conversion function.