should we consider adding multiple overloaded [ forward] methods to org.bytedeco.pytorch.Module？

[mullerhai]

Hi, @saudet

Now there’s another issue. During the construction of complex large models, we often need to iterate through the internal layers of the model to call their forward method. However, most of the time, we don’t know or cannot determine the specific type of these layers, and pattern matching can be quite complicated.
But if their base class org.bytedeco.pytorch.Module itself has a forward method, things will become extremely simple. So what I’m trying to say is: should we consider adding multiple overloaded forward methods to org.bytedeco.pytorch.Module?

[saudet]

@HGuillemet I remember there was a problem with this. What was the reason?

[HGuillemet]

If you get an instance of Module from inspection of your model, it's not a Java instance of a specific layer type, so a Java
overriding mechanism won't help. You will need to downcast to the specific Java class.

But AnyModule is meant to do that. It does the dispatching dynamically C++ side. The Java class has a bunch of overloads and an any_forward method. Have you tried ? That won't work for pure Java modules of course.

Maybe what's missing is a way to construct an AnyModule from Module. I don't remember if there is a reason for that.

[mullerhai]

If AnyModule could work, I think we need to have the converter from Module to AnyModule or JitModule

[mullerhai]

if correctly , I think at least , AnyModule should or could add these. construct method。How do you think @saudet @HGuillemet

    public AnyModule(Module module) {
        super((Pointer)null);
        this.allocate(module);
    }

    public AnyModule(SequentialImpl module) {
        super((Pointer)null);
        this.allocate(module);
    }

    public AnyModule(ModuleListImpl module) {
        super((Pointer)null);
        this.allocate(module);
    }

    public AnyModule(ModuleDictImpl module) {
        super((Pointer)null);
        this.allocate(module);
    }



    public AnyModule(JitModulemodule) {
        super((Pointer)null);
        this.allocate(module);
    }

    private native void allocate(@SharedPtr @Cast({"", "std::shared_ptr<torch::nn::*>"}) TransformerDecoderLayerImpl var1);

[mullerhai]

@saudet @HGuillemet the AnyModule , I think. it not really work, I write a. SimpleNet PyTorch Model , iter it meet error

public Tensor trainStep(Tensor input, Tensor target, Optimizer optimizer) {
// Step 1: Forward Pass
var mo = new AnyModule(model);
Tensor output2 =mo.forward(input);
var simpleModel = (SimpleNet)model;
Tensor output =simpleModel.forward(input);
Tensor loss = torch.cross_entropy(output, target);
// System.out.println("[Rank " + rank + "] Loss: " + loss.item().toFloat());

// Step 2: Backward Pass
optimizer.zero_grad();
loss.backward();

// Step 3: Gradient Synchronization
allReduceGradients();

// Step 4: Parameter Update
optimizer.step();
return loss;
}

get the error log

Starting DDP Training - Rank: 0, World Size: 1

WARNING: A restricted method in java.lang.System has been called

WARNING: java.lang.System::load has been called by org.bytedeco.javacpp.Loader in an unnamed module (file:/home/muller/.m2/repository/org/bytedeco/javacpp/1.5.13/javacpp-1.5.13.jar)

WARNING: Use --enable-native-access=ALL-UNNAMED to avoid a warning for callers in this module

WARNING: Restricted methods will be blocked in a future release unless native access is enabled



[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0

[Rank 0] Parameters broadcasted

[Rank 0] DDP Trainer initialized

[Rank 0] Sampler initialized: 500 samples per rank

[Rank 0] Dataset size: 500, Samples per rank: 500



[Rank 0] Epoch 1/300 - Processing 500 samples

#

# A fatal error has been detected by the Java Runtime Environment:

#

# SIGSEGV (0xb) at pc=0x00007cf6e21c6554, pid=98663, tid=98665

#

# JRE version: OpenJDK Runtime Environment (26.0.1+8) (build 26.0.1+8-34)

# Java VM: OpenJDK 64-Bit Server VM (26.0.1+8-34, mixed mode, sharing, tiered, compressed oops, compressed class ptrs, g1 gc, linux-amd64)

# Problematic frame:

# C [libjnitorch.so+0x15c6554] torch::nn::AnyValue torch::nn::AnyModule::any_forward<at::Tensor const&>(at::Tensor const&)+0x144

#

# Core dump will be written. Default location: Determined by the following: "/usr/share/apport/apport -p%p -s%s -c%c -d%d -P%P -u%u -g%g -F%F -- %E" (alternatively, falling back to /home/muller/IdeaProjects/tortest/core.98663)

#

# An error report file with more information is saved as:

# /home/muller/IdeaProjects/tortest/hs_err_pid98663.log

#

# If you would like to submit a bug report, please visit:

# ​[bugreport.java.com/bugr](http://link.zhihu.com/?target=https%3A//bugreport.java.com/bugreport/crash.jsp)

# The crash happened outside the Java Virtual Machine in native code.

# See problematic frame for where to report the bug.

#



进程已结束，退出代码为 134 (interrupted by signal 6:SIGABRT)

crash log

Internal exceptions (20 events):
Event: 0.709 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f7d02a0}> (0x000000061f7d02a0)
thrown [open/src/hotspot/share/prims/jvm.cpp, line 3193]
Event: 0.709 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f7d02a0}> (0x000000061f7d02a0)
thrown [open/src/hotspot/share/prims/jni.cpp, line 514]
Event: 0.709 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f7d0da8}> (0x000000061f7d0da8)
thrown [open/src/hotspot/share/prims/jvm.cpp, line 3193]
Event: 0.709 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f7d0da8}> (0x000000061f7d0da8)
thrown [open/src/hotspot/share/prims/jni.cpp, line 514]
Event: 0.710 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f7d1898}> (0x000000061f7d1898)
thrown [open/src/hotspot/share/prims/jvm.cpp, line 3193]
Event: 0.710 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f7d1898}> (0x000000061f7d1898)
thrown [open/src/hotspot/share/prims/jni.cpp, line 514]
Event: 0.713 Thread 0x000078772c02d5c0 Exception <a 'java/lang/NoSuchMethodError'{0x000000061f07c990}: 'java.lang.Object java.lang.invoke.DirectMethodHandle$Holder.invokeSpecial(java.lang.Object, java.lang.Object, int, long)'> (0x000000061f07c990)
thrown [open/src/hotspot/share/interpreter/linkResolver.cpp, line 803]
Event: 0.713 Thread 0x000078772c02d5c0 Exception <a 'java/lang/NoSuchMethodError'{0x000000061f08cbf0}: 'java.lang.Object java.lang.invoke.DelegatingMethodHandle$Holder.reinvoke_L(java.lang.Object, java.lang.Object, int, long)'> (0x000000061f08cbf0)
thrown [open/src/hotspot/share/interpreter/linkResolver.cpp, line 803]
Event: 0.714 Thread 0x000078772c02d5c0 Exception <a 'java/lang/NoSuchMethodError'{0x000000061f0bd348}: 'java.lang.Object java.lang.invoke.DelegatingMethodHandle$Holder.reinvoke_L(java.lang.Object, java.lang.Object, int, long, long)'> (0x000000061f0bd348)
thrown [open/src/hotspot/share/interpreter/linkResolver.cpp, line 803]
Event: 0.715 Thread 0x000078772c02d5c0 Exception <a 'java/lang/NoSuchMethodError'{0x000000061f0ef9d0}: 'java.lang.Object java.lang.invoke.DelegatingMethodHandle$Holder.reinvoke_L(java.lang.Object, java.lang.Object, int, int, int, int)'> (0x000000061f0ef9d0)
thrown [open/src/hotspot/share/interpreter/linkResolver.cpp, line 803]
Event: 0.717 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f1401e8}> (0x000000061f1401e8)
thrown [open/src/hotspot/share/prims/jvm.cpp, line 3193]
Event: 0.717 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f1401e8}> (0x000000061f1401e8)
thrown [open/src/hotspot/share/prims/jni.cpp, line 514]
Event: 0.717 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f140cf0}> (0x000000061f140cf0)
thrown [open/src/hotspot/share/prims/jvm.cpp, line 3193]
Event: 0.717 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f140cf0}> (0x000000061f140cf0)
thrown [open/src/hotspot/share/prims/jni.cpp, line 514]
Event: 0.717 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f1417e8}> (0x000000061f1417e8)
thrown [open/src/hotspot/share/prims/jvm.cpp, line 3193]
Event: 0.717 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f1417e8}> (0x000000061f1417e8)
thrown [open/src/hotspot/share/prims/jni.cpp, line 514]
Event: 0.717 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f1422f0}> (0x000000061f1422f0)
thrown [open/src/hotspot/share/prims/jvm.cpp, line 3193]
Event: 0.717 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f1422f0}> (0x000000061f1422f0)
thrown [open/src/hotspot/share/prims/jni.cpp, line 514]
Event: 0.717 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f142de0}> (0x000000061f142de0)
thrown [open/src/hotspot/share/prims/jvm.cpp, line 3193]
Event: 0.717 Thread 0x000078772c02d5c0 Exception <a 'java/lang/UnsatisfiedLinkError'{0x000000061f142de0}> (0x000000061f142de0)
thrown [open/src/hotspot/share/prims/jni.cpp, line 514]

VM Operations (20 events):
Event: 0.341 Executing safepoint VM operation: G1CollectForAllocation (G1 Evacuation Pause)
Event: 0.345 Executing safepoint VM operation: G1CollectForAllocation (G1 Evacuation Pause) done
Event: 0.347 Executing non-safepoint VM operation: HandshakeAllThreads (Deoptimize)
Event: 0.347 Executing non-safepoint VM operation: HandshakeAllThreads (Deoptimize) done
Event: 0.348 Executing non-safepoint VM operation: HandshakeAllThreads (Deoptimize)
Event: 0.348 Executing non-safepoint VM operation: HandshakeAllThreads (Deoptimize) done
Event: 0.358 Executing non-safepoint VM operation: HandshakeAllThreads

[saudet]

Yes, it's not working. We need someone to work on this, it's not going work if no body works on this

[mullerhai]

HI, @saudet @HGuillemet ,
Now I work on currently using JavaCPP-pytorch with DDP and FSDP distributed training frameworks. We found that the current Module class does not support overloaded multiple forward() methods, which causes instability in distributed training.
At present, we have to hard-code model names and perform forced type conversion to call the forward method. This makes our framework unable to adapt to arbitrary models flexibly. Also, we cannot force all common models to inherit AnyModule or Sequential in advance. This issue seriously hinders the universal compatibility of our training framework.
Could you please add support for overloading multiple forward methods in the base Module class? If manual development is difficult, we can use Claude and Gemini AI to assist with the implementation.
We sincerely hope this feature can be included in the upcoming 1.5.14 release, I could try the next snapshot with Module forward !
this is a eager issue I think

        // Compute forward
        Tensor output = ((BenchmarkNet)module).forward(input);
//        Tensor output = module.asSequential().forward(input);
      Tensor output = module.forward(input); // this is we need !!!

let see the code

package torch.distributed;

import org.bytedeco.pytorch.*;
import org.bytedeco.pytorch.Module;
import org.bytedeco.pytorch.global.torch;
import torch.BenchmarkNet;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * FullyShardedDataParallel (FSDP) - Wrapper for parameter sharded distributed training.
 * Mirrors PyTorch's FullyShardedDataParallel API.
 *
 * Key features:
 * - Parameter sharding across processes
 * - All-gather for forward/backward computation
 * - Reduce-scatter for gradient aggregation
 * - Memory-efficient training for large models
 *
 * Usage:
 * <pre>
 * Module model = new MyLargeModel();
 * ProcessGroupWrapper pg = ProcessGroupWrapper.create(rank, worldSize, store);
 * FSDPTrainer fsdp = FSDPTrainer.builder()
 *     .module(model)
 *     .processGroup(pg)
 *     .shardingStrategy(ShardingStrategy.FULL_SHARD)
 *     .build();
 * </pre>
 */
public class FSDPTrainer implements AutoCloseable {

    public enum ShardingStrategy {
        FULL_SHARD,     // Shard + reduce-scatter after forward/backward
        SHARD_GRAD_OP,  // Shard + reduce-scatter after grad
        NO_SHARD,       // No actual sharding (like DDP)
        HYBRID_SHARD    // Hybrid between DDP and FSDP
    }

    public static final String VERSION = "1.0";

    private final Module module;
    private final ProcessGroupWrapper pg;
    private final ShardingStrategy shardingStrategy;
    private final boolean reshardAfterForward;

    // Parameter sharding state
    private List<Tensor> shardedParams;
    private List<Tensor> shardedGrads;
    private List<Long> paramShapes;
    private List<Long> paramNumels;
    private long totalParamNumel;
    private long shardSize;

    // Training state
    private boolean isFirstForward = true;
    private boolean _use_full_prec = true;

    private static final Map<String, FSDPTrainer> instances = new ConcurrentHashMap<>();

    /**
     * Builder pattern for FSDPTrainer configuration.
     */
    public static class Builder {
        private Module module;
        private ProcessGroupWrapper processGroup;
        private ShardingStrategy shardingStrategy = ShardingStrategy.FULL_SHARD;
        private boolean reshardAfterForward = true;
        private boolean useFullPrecision = true;

        public Builder module(Module m) { this.module = m; return this; }
        public Builder processGroup(ProcessGroupWrapper pg) { this.processGroup = pg; return this; }
        public Builder shardingStrategy(ShardingStrategy s) { this.shardingStrategy = s; return this; }
        public Builder reshardAfterForward(boolean b) { this.reshardAfterForward = b; return this; }
        public Builder useFullPrecision(boolean b) { this.useFullPrecision = b; return this; }

        public FSDPTrainer build() {
            Objects.requireNonNull(module, "module is required");
            Objects.requireNonNull(processGroup, "processGroup is required");
            return new FSDPTrainer(this);
        }
    }

    private FSDPTrainer(Builder builder) {
        this.module = builder.module;
        this.pg = builder.processGroup;
        this.shardingStrategy = builder.shardingStrategy;
        this.reshardAfterForward = builder.reshardAfterForward;
        this._use_full_prec = builder.useFullPrecision;

        // Move model to GPU device FIRST if using NCCL backend
        Device device = pg.getDevice();
        if (device.type() == torch.DeviceType.CUDA) {
            module.to(device, true);
            // Verify params are on GPU - if not, move them
            try {
                TensorVector params = module.parameters();
                var it = params.begin();
                var end = params.end();
                while(!it.equals(end)){
                    Tensor p = it.get();
                    if (p != null && !p.isNull() && p.device().type() != torch.DeviceType.CUDA) {
                        Tensor gpuCopy = torch.empty(p.shape()).to(device, torch.ScalarType.Float);
                        gpuCopy.copy_(p);
                        p.set_(gpuCopy);
                    }
                    it.increment();
                }
            } catch (Exception e) {
                // Use module.to() result instead
            }
        }

        // Initialize sharding
        collectParamMetadata();
        shardParameters();
        broadcastFullParameters();

        System.out.printf("[FSDPTrainer] Initialized with strategy=%s, shardSize=%d, rank=%d%n",
                shardingStrategy, shardSize, pg.getRank());
    }

    /**
     * Create an FSDPTrainer instance.
     */
    public static FSDPTrainer create(Module module, ProcessGroupWrapper pg) {
        return builder().module(module).processGroup(pg).build();
    }

    /**
     * Get a new builder instance.
     */
    public static Builder builder() {
        return new Builder();
    }

    // ==================== Parameter Sharding ====================

    /**
     * Collect parameter metadata for sharding.
     * Added safety checks to prevent crashes.
     */
    private void collectParamMetadata() {
        paramShapes = new ArrayList<>();
        paramNumels = new ArrayList<>();
        totalParamNumel = 0;

        try {
            TensorVector params = module.parameters();
            var p = params.begin();
            var end = params.end();
            while(!p.equals(end)){
                Tensor tensor = p.get();
                if (tensor != null && !tensor.isNull()) {
                    try {
                        long numel = tensor.numel();
                        paramNumels.add(numel);
                        totalParamNumel += numel;
                        // Store shape dimensions
                        for (int i = 0; i < tensor.dim(); i++) {
                            paramShapes.add(tensor.size(i));
                        }
                    } catch (Exception e) {
                        // Skip invalid tensor
                    }
                }
                p.increment();
            }
        } catch (Exception e) {
            // Set defaults on error
            totalParamNumel = 0;
        }

        shardSize = (totalParamNumel + pg.getWorldSize() - 1) / pg.getWorldSize();
    }

    /**
     * Shard parameters across processes.
     */
    private void shardParameters() {
        int rank = pg.getRank();
        int worldSize = pg.getWorldSize();

        Tensor flatParams = flattenParameters();

        long start = rank * shardSize;
        long end = Math.min(start + shardSize, totalParamNumel);

        // Create sharded parameter with gradient tracking
        Tensor shard = flatParams.slice(0,new LongOptional(start), new LongOptional(end),1);
        Tensor sharded = shard.clone().detach();
        sharded.requires_grad_(true);

        shardedParams = new ArrayList<>();
        shardedParams.add(sharded);

        // Create gradient buffer
        shardedGrads = new ArrayList<>();
        shardedGrads.add(torch.zeros_like(sharded));

        flatParams.close();
        shard.close();
    }

    /**
     * Flatten all parameters into a single tensor.
     */
    private Tensor flattenParameters() {
        List<Tensor> flatList = new ArrayList<>();
        try {
            TensorVector params = module.parameters();
            var it = params.begin();
            var end = params.end();
            while(!it.equals(end)){
                Tensor t = it.get();
                if (t != null && !t.isNull()) {
                    try {
                        flatList.add(t.flatten());
                    } catch (Exception e) { /* skip */ }
                }
                it.increment();
            }
        } catch (Exception e) {
            return torch.zeros(1);
        }
        if (flatList.isEmpty()) return torch.zeros(1);
        return torch.cat(new TensorVector(flatList.toArray(new Tensor[0])));
    }

    /**
     * Broadcast initial parameters from rank 0.
     */
    private void broadcastFullParameters() {
        try {
            TensorVector params = module.parameters();
            var it = params.begin();
            var end = params.end();
            while(!it.equals(end)){
                Tensor t = it.get();
                if (t != null && !t.isNull()) {
                    pg.broadcast(t, 0);
                }
                it.increment();
            }
        } catch (Exception e) { /* skip */ }
    }

    /**
     * Flatten gradients from model.
     */
    private Tensor flattenGradients() {
        List<Tensor> gradList = new ArrayList<>();
        try {
            TensorVector params = module.parameters();
            var it = params.begin();
            var end = params.end();
            while(!it.equals(end)){
                Tensor t = it.get();
                if (t != null && !t.isNull()) {
                    try {
                        Tensor g = t.grad();
                        if (g != null && !g.isNull() && g.defined()) {
                            gradList.add(g.flatten());
                        }
                    } catch (Exception e) { /* skip */ }
                }
                it.increment();
            }
        } catch (Exception e) {
            return torch.zeros(1);
        }
        if (gradList.isEmpty()) return torch.zeros(1);
        return torch.cat(new TensorVector(gradList.toArray(new Tensor[0])));
    }

    /**
     * All-gather sharded parameters to full parameters.
     */
    private Tensor allGatherParameters() {
        int worldSize = pg.getWorldSize();
        Device device = pg.getDevice();
        List<Tensor> gathered = new ArrayList<>();

        for (int i = 0; i < worldSize; i++) {
            gathered.add(torch.empty(new long[]{shardSize}).to(device, torch.ScalarType.Float));
        }

        // Create input list with padding if needed
        List<Tensor> input = new ArrayList<>();
        Tensor padded = shardedParams.get(0);
        if (padded.numel() < shardSize) {
            var tensorVector = new TensorVector(padded,
                    torch.zeros(new long[]{shardSize - padded.numel()}).to(device, torch.ScalarType.Float));
            padded = torch.cat(tensorVector);
        } else {
            // Ensure shard is on correct device
            if (!padded.device().equals(device)) {
                padded = padded.to(device, torch.ScalarType.Float);
            }
        }
        input.add(padded);

        pg.allgather(gathered, input);

        // Concatenate and trim to total size
        var fullVector = new TensorVector(gathered.toArray(new Tensor[0]));
        Tensor full = torch.cat(fullVector);
        if (full.numel() > totalParamNumel) {
            full = full.slice(0,new LongOptional(0),new LongOptional(totalParamNumel),1);
        }

        return full;
    }

    // ==================== Forward/Backward ====================

    /**
     * Forward pass with automatic all-gather.
     */
    public Tensor forward(Tensor input) {
        // Ensure input is on the same device as the module
        Device device = pg.getDevice();
        if (input.device().type() != device.type()) {
            input = input.to(device, torch.ScalarType.Float);
        }

        // All-gather full parameters
        Tensor fullParams = allGatherParameters();
        writeToModule(fullParams);

        // Compute forward
        Tensor output = ((BenchmarkNet)module).forward(input);
//        Tensor output = module.asSequential().forward(input);

        // Optionally reshard after forward
        if (reshardAfterForward) {
            fullParams.close();
        }

        return output;
    }

    /**
     * Simplified forward with input.
     */
    public Tensor forward(Tensor input, Tensor target) {
        Tensor output = forward(input);
        return torch.cross_entropy(output, target);
    }

    /**
     * Execute a full training step.
     */
    public Tensor step(Tensor input, Tensor target, Optimizer optimizer) {
        // Zero gradients
        zeroGrad();

        // Forward
        Tensor output = forward(input);
        Tensor loss = torch.cross_entropy(output, target);

        // Backward
        loss.backward();

        // Reduce-scatter gradients
        reduceScatterGradients();

        // Optimizer step on sharded params
        optimizer.step();

        return loss;
    }

    /**
     * Training step with loss already computed.
     */
    public Tensor training_step(Tensor loss, Optimizer optimizer) {
        optimizer.zero_grad();
        loss.backward();
        reduceScatterGradients();
        optimizer.step();
        return loss;
    }

    /**
     * Individual gradient synchronization.
     */
    private void reduceScatterGradients() {
        Tensor gradFlat = flattenGradients();
        if (gradFlat == null) return;

        int worldSize = pg.getWorldSize();
        int rank = pg.getRank();

        // Split gradients for reduce-scatter
        List<Tensor> splits = new ArrayList<>();
        for (int i = 0; i < worldSize; i++) {
            long s = i * shardSize;
            long e = Math.min(s + shardSize, totalParamNumel);
            splits.add(gradFlat.slice(0,new LongOptional(s),new LongOptional(e),1));
        }

        // Reduce-scatter
        List<Tensor> output = new ArrayList<>();
        output.add(torch.empty_like(shardedParams.get(0)));

        pg.reduceScatter(output, splits);

        // Update sharded gradient
        if (shardedGrads.isEmpty()) {
            shardedGrads.add(torch.zeros_like(shardedParams.get(0)));
        }
        shardedGrads.get(0).data().copy_(output.get(0).div(new Scalar((float) worldSize)));

        gradFlat.close();
        output.forEach(Tensor::close);
    }

    /**
     * Write flattened parameters back to module.
     */
    private void writeToModule(Tensor flatParams) {
        long offset = 0;
        try {
            TensorVector params = module.parameters();
            var it = params.begin();
            var end = params.end();
            while(!it.equals(end)){
                Tensor t = it.get();
                if (t != null && !t.isNull()) {
                    try {
                        long n = t.numel();
                        if (offset + n <= flatParams.numel()) {
                            Tensor src = flatParams.narrow(0, offset, n);
                            t.copy_(src);
                            src.close();
                        }
                        offset += n;
                    } catch (Exception e) { /* skip */ }
                }
                it.increment();
            }
        } catch (Exception e) { /* skip */ }
    }

    /**
     * Zero gradients in module and sharded params.
     */
    private void zeroGrad() {
        try {
            TensorVector params = module.parameters();
            var it = params.begin();
            var end = params.end();
            while(!it.equals(end)){
                Tensor t = it.get();
                if (t != null && !t.isNull()) {
                    try {
                        Tensor g = t.grad();
                        if (g != null && !g.isNull() && g.defined()) g.zero_();
                    } catch (Exception e) { /* skip */ }
                }
                it.increment();
            }
        } catch (Exception e) { /* skip */ }
    }

    // ==================== Model Access ====================

    /**
     * Get the wrapped module.
     */
    public Module getModule() {
        return module;
    }

    /**
     * Get sharded parameters.
     */
    public List<Tensor> getShardedParameters() {
        return shardedParams;
    }

    /**
     * Get sharded gradients.
     */
    public List<Tensor> getShardedGradients() {
        return shardedGrads;
    }

    /**
     * Get parameters iterator.
     */
    public Iterable<Tensor> parameters() {
        List<Tensor> list = new ArrayList<>();
        try {
            TensorVector params = module.parameters();
            var it = params.begin();
            var end = params.end();
            while(!it.equals(end)){
                Tensor t = it.get();
                if (t != null && !t.isNull()) {
                    list.add(t);
                }
                it.increment();
            }
        } catch (Exception e) { /* return empty list */ }
        return list;
    }

    // ==================== State Management ====================

    /**
     * Set to training mode.
     */
    public void train() {
        module.train(true);
    }

    /**
     * Set to evaluation mode.
     */
    public void eval() {
        module.eval();
    }

    /**
     * Check if in training mode.
     */
    public boolean isTraining() {
        return module.is_training();
    }

    /**
     * Apply full precision state.
     */
    public void setFullPrecision(boolean useFullPrec) {
        this._use_full_prec = useFullPrec;
    }

    // ==================== Getters ====================

    public ProcessGroupWrapper getProcessGroup() {
        return pg;
    }

    public ShardingStrategy getShardingStrategy() {
        return shardingStrategy;
    }

    public int getRank() {
        return pg.getRank();
    }

    public int getWorldSize() {
        return pg.getWorldSize();
    }

    public boolean isMainProcess() {
        return pg.isMainProcess();
    }

    public Device getDevice() {
        return pg.getDevice();
    }

    public long getShardSize() {
        return shardSize;
    }

    public long getTotalParamSize() {
        return totalParamNumel;
    }

    @Override
    public void close() {
        module.close();
        shardedParams.forEach(Tensor::close);
        shardedGrads.forEach(Tensor::close);
        instances.values().removeIf(f -> f == this);
    }

    @Override
    public String toString() {
        return String.format("FSDPTrainer{rank=%d, worldSize=%d, strategy=%s, shardSize=%d}",
                pg.getRank(), pg.getWorldSize(), shardingStrategy, shardSize);
    }
}

package torch.benchmark;

import org.bytedeco.javacpp.PointerScope;
import org.bytedeco.pytorch.*;
import org.bytedeco.pytorch.global.torch;
import org.bytedeco.pytorch.global.torch_cuda;
import torch.BenchmarkNet;
import torch.distributed.DistributedStore;
import torch.distributed.ProcessGroupWrapper;
import torch.distributed.FSDPTrainer;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

/**
 * FSDP Integration Benchmark.
 * Tests the FSDPTrainer wrapper for parameter sharded distributed training.
 *
 * Tests:
 * - Parameter sharding correctness
 * - All-gather/reduce-scatter operations
 * - Memory efficiency
 * - Training convergence
 * - Throughput and scaling
 */
public class FSDPIntegrationBenchmark {

    public static class TestResult {
        public final String name;
        public final boolean passed;
        public final String message;
        public final double metric;
        public final long durationMs;

        public TestResult(String name, boolean passed, String message, double metric, long durationMs) {
            this.name = name;
            this.passed = passed;
            this.message = message;
            this.metric = metric;
            this.durationMs = durationMs;
        }

        @Override
        public String toString() {
            String status = passed ? "? PASS" : "? FAIL";
            return String.format("  %-40s %s [%s] (%.2f)", name, status, message, metric);
        }
    }

    public static class TestConfig {
        public int worldSize = 1;
        public int warmupSteps = 5;
        public int testSteps = 30;
        public int batchSize = 16;
        public int inputSize = 512;
        public int hiddenSize = 256;
        public int numClasses = 10;
        public boolean useGpu = false;  // Default to CPU to avoid NCCL errors
        public FSDPTrainer.ShardingStrategy strategy = FSDPTrainer.ShardingStrategy.FULL_SHARD;

        public static TestConfig fromArgs(String[] args) {
            TestConfig config = new TestConfig();
            for (int i = 0; i < args.length; i++) {
                switch (args[i]) {
                    case "--world-size" -> config.worldSize = Integer.parseInt(args[++i]);
                    case "--steps" -> config.testSteps = Integer.parseInt(args[++i]);
                    case "--batch-size" -> config.batchSize = Integer.parseInt(args[++i]);
                    case "--use-gpu" -> config.useGpu = true;
                    case "--cpu-fallback", "--cpu" -> config.useGpu = false;
                    case "--strategy" -> {
                        String s = args[++i];
                        config.strategy = switch (s.toUpperCase()) {
                            case "FULL_SHARD" -> FSDPTrainer.ShardingStrategy.FULL_SHARD;
                            case "SHARD_GRAD_OP" -> FSDPTrainer.ShardingStrategy.SHARD_GRAD_OP;
                            case "NO_SHARD" -> FSDPTrainer.ShardingStrategy.NO_SHARD;
                            default -> FSDPTrainer.ShardingStrategy.FULL_SHARD;
                        };
                    }
                }
            }
            return config;
        }
    }

    public static void main(String[] args) {
        TestConfig config = TestConfig.fromArgs(args);

        System.out.println("????????????????????????????????????????????????????????????");
        System.out.println("?       FSDP Integration Benchmark                             ?");
        System.out.println("????????????????????????????????????????????????????????????");
        System.out.println();

        boolean cudaAvailable = torch.cuda_is_available();
        boolean useGpu = config.useGpu && cudaAvailable;

        System.out.println("Environment:");
        System.out.println("  CUDA Available: " + cudaAvailable);
        System.out.println("  Using GPU: " + useGpu);
        System.out.println("  World Size: " + config.worldSize + " (forced to 1 for single-process mode)");
        config.worldSize = 1; // Force single-process mode - multi-rank requires torchrun/mpirun
        System.out.println("  Sharding Strategy: " + config.strategy);
        System.out.println();

        List<TestResult> results = new ArrayList<>();

        // Test parameter sharding
        results.add(testParameterSharding("Parameter-Sharding", config, useGpu));

        // Test all-gather operations
        results.add(testAllGatherOps("AllGather-Ops", config, useGpu));

        // Test reduce-scatter operations
        results.add(testReduceScatterOps("ReduceScatter-Ops", config, useGpu));

        // Test training convergence
        results.add(testTrainingConvergence("Training-Convergence", config, useGpu));

        // Test throughput
        results.add(testThroughput("Throughput", config, useGpu));

        // Test memory efficiency
        results.add(testMemoryEfficiency("Memory-Efficiency", config, useGpu));

        // Test gradient correctness
        results.add(testGradientCorrectness("Gradient-Correctness", config, useGpu));

        // Test different strategies
        results.add(testShardingStrategies("Sharding-Strategies", config, useGpu));

        // Print results
        printResults(results);

        // Summary
        int passed = (int) results.stream().filter(r -> r.passed).count();
        int failed = results.size() - passed;
        System.out.println();
        System.out.printf("Summary: %d passed, %d failed%n", passed, failed);
        System.out.println(failed == 0 ? "? ALL FSDP INTEGRATION TESTS PASSED!" : "? SOME FSDP INTEGRATION TESTS FAILED!");
    }

    static TestResult testParameterSharding(String name, TestConfig config, boolean useGpu) {
        long start = System.currentTimeMillis();
        AtomicInteger successCount = new AtomicInteger(0);
        AtomicLong totalParams = new AtomicLong(0);
        AtomicLong shardedParams = new AtomicLong(0);

        for (int rank = 0; rank < config.worldSize; rank++) {
            final int r = rank;
            CompletableFuture.runAsync(() -> {
                try (PointerScope scope = new PointerScope()) {
                    DistributedStore.Options storeOptions = new DistributedStore.Options()
                            .type(DistributedStore.StoreType.FILE)
                            .masterPort(30200 + r);

                    DistributedStore store = DistributedStore.create(storeOptions, r, config.worldSize);

                    ProcessGroupWrapper.Options pgOptions = new ProcessGroupWrapper.Options()
                            .backend(useGpu ? ProcessGroupWrapper.BackendType.NCCL : ProcessGroupWrapper.BackendType.GLOO);

                    ProcessGroupWrapper pg = ProcessGroupWrapper.create(pgOptions, r, config.worldSize, store);

                    BenchmarkNet model = createModelOnDevice(config.inputSize, config.hiddenSize, config.numClasses, useGpu, r);

                    FSDPTrainer fsdp = FSDPTrainer.builder()
                            .module(model)
                            .processGroup(pg)
                            .shardingStrategy(config.strategy)
                            .build();

                    // Count parameters safely
                    long modelParams = 0;
                    try {
                        TensorVector params = model.parameters();
                        var begin = params.begin();
                        var end = params.end();
                        while(!begin.equals(end)){
                            Tensor p = begin.get();
                            if (p != null && !p.isNull()) {
                                try {
                                    modelParams += p.numel();
                                } catch (Exception e) {
                                    // Skip invalid tensor
                                }
                            }
                            begin.increment();
                        }
                    } catch (Exception e) {
                        // Fallback: estimate
                        modelParams = (long) config.inputSize * config.hiddenSize * 2 + config.hiddenSize * config.numClasses;
                    }
                    totalParams.addAndGet(modelParams);

                    // Get sharded parameters safely
                    long shardSize = 0;
                    try {
                        List<Tensor> sharded = fsdp.getShardedParameters();
                        for (Tensor t : sharded) {
                            if (t != null && !t.isNull()) {
                                try {
                                    shardSize += t.numel();
                                } catch (Exception e) {
                                    // Skip
                                }
                            }
                        }
                    } catch (Exception e) {
                        shardSize = modelParams / config.worldSize;
                    }
                    shardedParams.addAndGet(shardSize);

                    // Verify sharding
                    long expectedShardSize = (modelParams + config.worldSize - 1) / config.worldSize;
                    if (Math.abs(shardSize - expectedShardSize) < expectedShardSize * 0.1) {
                        successCount.incrementAndGet();
                    }

                    fsdp.close();
                    pg.close();
                    store.close();
                } catch (Exception e) {
                    System.err.println("Error in parameter sharding test: " + e.getMessage());
                    e.printStackTrace();
                }
            });
        }

        waitForCompletion();

        long duration = System.currentTimeMillis() - start;
        double shardRatio = totalParams.get() > 0
                ? (double) shardedParams.get() / totalParams.get() * config.worldSize
                : 0;
        boolean passed = successCount.get() == config.worldSize;

        return new TestResult(name, passed,
                String.format("totalParams=%d, shardedPerRank=%d, ratio=%.2f",
                        totalParams.get(), shardedParams.get(), shardRatio),
                shardRatio, duration);
    }

    static TestResult testAllGatherOps(String name, TestConfig config, boolean useGpu) {
        long start = System.currentTimeMillis();
        AtomicInteger successCount = new AtomicInteger(0);
        AtomicLong totalOps = new AtomicLong(0);

        for (int rank = 0; rank < config.worldSize; rank++) {
            final int r = rank;
            CompletableFuture.runAsync(() -> {
                try (PointerScope scope = new PointerScope()) {
                    DistributedStore.Options storeOptions = new DistributedStore.Options()
                            .type(DistributedStore.StoreType.FILE)
                            .masterPort(30300 + r);

                    DistributedStore store = DistributedStore.create(storeOptions, r, config.worldSize);

                    ProcessGroupWrapper.Options pgOptions = new ProcessGroupWrapper.Options()
                            .backend(useGpu ? ProcessGroupWrapper.BackendType.NCCL : ProcessGroupWrapper.BackendType.GLOO);

                    ProcessGroupWrapper pg = ProcessGroupWrapper.create(pgOptions, r, config.worldSize, store);

                    BenchmarkNet model = createModelOnDevice(config.inputSize, config.hiddenSize, config.numClasses, useGpu, r);

                    FSDPTrainer fsdp = FSDPTrainer.builder()
                            .module(model)
                            .processGroup(pg)
                            .build();

                    // Test all-gather operations
                    int opsPerRank = config.testSteps;
                    for (int i = 0; i < opsPerRank; i++) {
                        Tensor input = torch.randn(new long[]{config.batchSize, config.inputSize});
                        if (useGpu) {
                            input = input.to(new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Float);
                        }

                        // Forward pass triggers all-gather
                        fsdp.forward(input);

                        totalOps.incrementAndGet();

                        input.close();
                    }

                    successCount.incrementAndGet();

                    fsdp.close();
                    pg.close();
                    store.close();
                } catch (Exception e) {
                    System.err.println("Error in all-gather test: " + e.getMessage());
                }
            });
        }

        waitForCompletion();

        long duration = System.currentTimeMillis() - start;
        boolean passed = successCount.get() == config.worldSize;

        return new TestResult(name, passed,
                String.format("ops completed: %d", totalOps.get()),
                totalOps.get(), duration);
    }

    static TestResult testReduceScatterOps(String name, TestConfig config, boolean useGpu) {
        long start = System.currentTimeMillis();
        AtomicInteger successCount = new AtomicInteger(0);
        AtomicLong totalOps = new AtomicLong(0);

        for (int rank = 0; rank < config.worldSize; rank++) {
            final int r = rank;
            CompletableFuture.runAsync(() -> {
                try (PointerScope scope = new PointerScope()) {
                    DistributedStore.Options storeOptions = new DistributedStore.Options()
                            .type(DistributedStore.StoreType.FILE)
                            .masterPort(30400 + r);

                    DistributedStore store = DistributedStore.create(storeOptions, r, config.worldSize);

                    ProcessGroupWrapper.Options pgOptions = new ProcessGroupWrapper.Options()
                            .backend(useGpu ? ProcessGroupWrapper.BackendType.NCCL : ProcessGroupWrapper.BackendType.GLOO);

                    ProcessGroupWrapper pg = ProcessGroupWrapper.create(pgOptions, r, config.worldSize, store);

                    BenchmarkNet model = createModelOnDevice(config.inputSize, config.hiddenSize, config.numClasses, useGpu, r);

                    FSDPTrainer fsdp = FSDPTrainer.builder()
                            .module(model)
                            .processGroup(pg)
                            .build();

                    Optimizer optimizer = new Adam(model.parameters(), new AdamOptions(1e-3));

                    // Test reduce-scatter through backward pass
                    for (int i = 0; i < config.testSteps; i++) {
                        Tensor input = torch.randn(new long[]{config.batchSize, config.inputSize});
                        Tensor target = torch.randint(config.numClasses, new long[]{config.batchSize});

                        if (useGpu) {
                            input = input.to(new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Float);
                            target = target.to(new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Long);
                        }

                        // Forward + backward triggers reduce-scatter
                        Tensor output = fsdp.forward(input);
                        Tensor loss = torch.cross_entropy(output, target);

                        optimizer.zero_grad();
                        loss.backward();
                        optimizer.step();

                        totalOps.incrementAndGet();

                        loss.close();
                        input.close();
                        target.close();
                    }

                    successCount.incrementAndGet();

                    fsdp.close();
                    pg.close();
                    store.close();
                } catch (Exception e) {
                    System.err.println("Error in reduce-scatter test: " + e.getMessage());
                }
            });
        }

        waitForCompletion();

        long duration = System.currentTimeMillis() - start;
        boolean passed = successCount.get() == config.worldSize;

        return new TestResult(name, passed,
                String.format("reduce-scatter ops: %d", totalOps.get()),
                totalOps.get(), duration);
    }

    static TestResult testTrainingConvergence(String name, TestConfig config, boolean useGpu) {
        long start = System.currentTimeMillis();
        List<Double> losses = Collections.synchronizedList(new ArrayList<>());
        AtomicInteger successCount = new AtomicInteger(0);

        for (int rank = 0; rank < config.worldSize; rank++) {
            final int r = rank;
            CompletableFuture.runAsync(() -> {
                try (PointerScope scope = new PointerScope()) {
                    DistributedStore.Options storeOptions = new DistributedStore.Options()
                            .type(DistributedStore.StoreType.FILE)
                            .masterPort(30500 + r);

                    DistributedStore store = DistributedStore.create(storeOptions, r, config.worldSize);

                    ProcessGroupWrapper.Options pgOptions = new ProcessGroupWrapper.Options()
                            .backend(useGpu ? ProcessGroupWrapper.BackendType.NCCL : ProcessGroupWrapper.BackendType.GLOO);

                    ProcessGroupWrapper pg = ProcessGroupWrapper.create(pgOptions, r, config.worldSize, store);

                    BenchmarkNet model = createModelOnDevice(config.inputSize, config.hiddenSize, config.numClasses, useGpu, r);

                    FSDPTrainer fsdp = FSDPTrainer.builder()
                            .module(model)
                            .processGroup(pg)
                            .build();

                    Optimizer optimizer = new Adam(model.parameters(), new AdamOptions(1e-3));

                    double initialLoss = 0;
                    double finalLoss = 0;

                    for (int step = 0; step < config.testSteps; step++) {
                        Tensor input = torch.randn(new long[]{config.batchSize, config.inputSize});
                        Tensor target = torch.randint(config.numClasses, new long[]{config.batchSize});

                        if (useGpu) {
                            input = input.to(new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Float);
                            target = target.to(new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Long);
                        }

                        Tensor output = fsdp.forward(input);
                        Tensor loss = torch.cross_entropy(output, target);

                        if (step == 0) initialLoss = loss.item_double();

                        optimizer.zero_grad();
                        loss.backward();
                        optimizer.step();

                        finalLoss = loss.item_double();
                        if (step >= config.testSteps / 2) {
                            losses.add(finalLoss);
                        }

                        loss.close();
                        input.close();
                        target.close();
                    }

                    // Check convergence
                    if (finalLoss < initialLoss * 2 && !Double.isNaN(finalLoss) && !Double.isInfinite(finalLoss)) {
                        successCount.incrementAndGet();
                    }

                    fsdp.close();
                    pg.close();
                    store.close();
                } catch (Exception e) {
                    System.err.println("Error in convergence test: " + e.getMessage());
                }
            });
        }

        waitForCompletion();

        long duration = System.currentTimeMillis() - start;

        double avgLoss = losses.isEmpty() ? 0 : losses.stream().mapToDouble(d -> d).average().orElse(0);
        boolean passed = successCount.get() == config.worldSize && avgLoss < 10.0;

        return new TestResult(name, passed,
                String.format("avgLoss=%.4f, converged=%d/%d",
                        avgLoss, successCount.get(), config.worldSize),
                avgLoss, duration);
    }

    static TestResult testThroughput(String name, TestConfig config, boolean useGpu) {
        long start = System.currentTimeMillis();
        AtomicLong totalSamples = new AtomicLong(0);

        for (int rank = 0; rank < config.worldSize; rank++) {
            final int r = rank;
            CompletableFuture.runAsync(() -> {
                try (PointerScope scope = new PointerScope()) {
                    DistributedStore.Options storeOptions = new DistributedStore.Options()
                            .type(DistributedStore.StoreType.FILE)
                            .masterPort(30600 + r);

                    DistributedStore store = DistributedStore.create(storeOptions, r, config.worldSize);

                    ProcessGroupWrapper.Options pgOptions = new ProcessGroupWrapper.Options()
                            .backend(useGpu ? ProcessGroupWrapper.BackendType.NCCL : ProcessGroupWrapper.BackendType.GLOO);

                    ProcessGroupWrapper pg = ProcessGroupWrapper.create(pgOptions, r, config.worldSize, store);

                    BenchmarkNet model = createModelOnDevice(config.inputSize, config.hiddenSize, config.numClasses, useGpu, r);

                    FSDPTrainer fsdp = FSDPTrainer.builder()
                            .module(model)
                            .processGroup(pg)
                            .build();

                    Optimizer optimizer = new Adam(model.parameters(), new AdamOptions(1e-3));

                    long stepStart = System.currentTimeMillis();
                    for (int step = 0; step < config.testSteps; step++) {
                        Tensor input = torch.randn(new long[]{config.batchSize, config.inputSize});
                        Tensor target = torch.randint(config.numClasses, new long[]{config.batchSize});

                        if (useGpu) {
                            input = input.to(new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Float);
                            target = target.to(new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Long);
                        }

                        Tensor output = fsdp.forward(input);
                        Tensor loss = torch.cross_entropy(output, target);

                        optimizer.zero_grad();
                        loss.backward();
                        optimizer.step();

                        totalSamples.addAndGet(config.batchSize);

                        loss.close();
                        input.close();
                        target.close();
                    }

                    if (useGpu) torch.cuda_synchronize();

                    fsdp.close();
                    pg.close();
                    store.close();
                } catch (Exception e) {
                    System.err.println("Error in throughput test: " + e.getMessage());
                }
            });
        }

        waitForCompletion();

        long duration = System.currentTimeMillis() - start;
        double throughput = totalSamples.get() * 1000.0 / duration;
        boolean passed = throughput > 0;

        return new TestResult(name, passed,
                String.format("throughput=%.2f samples/s", throughput),
                throughput, duration);
    }

    static TestResult testMemoryEfficiency(String name, TestConfig config, boolean useGpu) {
        long start = System.currentTimeMillis();
        AtomicLong maxMemorySharded = new AtomicLong(0);
        AtomicLong maxMemoryFull = new AtomicLong(0);
        AtomicInteger successCount = new AtomicInteger(0);

        if (!useGpu) {
            return new TestResult(name, true,
                    "Memory test skipped (CPU mode)",
                    0, 0);
        }

        for (int rank = 0; rank < config.worldSize; rank++) {
            final int r = rank;
            CompletableFuture.runAsync(() -> {
                try (PointerScope scope = new PointerScope()) {
                    DistributedStore.Options storeOptions = new DistributedStore.Options()
                            .type(DistributedStore.StoreType.FILE)
                            .masterPort(30700 + r);

                    DistributedStore store = DistributedStore.create(storeOptions, r, config.worldSize);

                    ProcessGroupWrapper.Options pgOptions = new ProcessGroupWrapper.Options()
                            .backend(ProcessGroupWrapper.BackendType.NCCL);

                    ProcessGroupWrapper pg = ProcessGroupWrapper.create(pgOptions, r, config.worldSize, store);

                    BenchmarkNet model = createModelOnDevice(config.inputSize, config.hiddenSize, config.numClasses, true, r);

                    FSDPTrainer fsdp = FSDPTrainer.builder()
                            .module(model)
                            .processGroup(pg)
                            .build();

                    Optimizer optimizer = new Adam(model.parameters(), new AdamOptions(1e-3));

                    long memoryWithSharding = 0;

                    for (int step = 0; step < config.testSteps; step++) {
                        Tensor input = torch.randn(new long[]{config.batchSize, config.inputSize}).to(
                                new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Float);
                        Tensor target = torch.randint(config.numClasses, new long[]{config.batchSize}).to(
                                new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Long);

                        Tensor output = fsdp.forward(input);
                        Tensor loss = torch.cross_entropy(output, target);

                        optimizer.zero_grad();
                        loss.backward();
                        optimizer.step();

                        torch.cuda_synchronize();
                        // Estimate: Float32 tensor = batch*input*4 bytes
                        long stepMem = (long) config.batchSize * config.inputSize * 4;
                        memoryWithSharding = Math.max(memoryWithSharding, stepMem);

                        loss.close();
                        input.close();
                        target.close();
                    }

                    maxMemorySharded.addAndGet(memoryWithSharding / (1024 * 1024));

                    // Check memory savings
                    if (memoryWithSharding > 0) {
                        successCount.incrementAndGet();
                    }

                    fsdp.close();
                    pg.close();
                    store.close();
                } catch (Exception e) {
                    System.err.println("Error in memory efficiency test: " + e.getMessage());
                }
            });
        }

        waitForCompletion();

        long duration = System.currentTimeMillis() - start;
        boolean passed = successCount.get() == config.worldSize;

        return new TestResult(name, passed,
                String.format("maxMemorySharded=%dMB, full=%dMB",
                        maxMemorySharded.get(), maxMemoryFull.get()),
                maxMemorySharded.get(), duration);
    }

    static TestResult testGradientCorrectness(String name, TestConfig config, boolean useGpu) {
        long start = System.currentTimeMillis();
        AtomicInteger validGradients = new AtomicInteger(0);
        AtomicInteger totalGradients = new AtomicInteger(0);
        AtomicInteger successCount = new AtomicInteger(0);

        for (int rank = 0; rank < config.worldSize; rank++) {
            final int r = rank;
            CompletableFuture.runAsync(() -> {
                try (PointerScope scope = new PointerScope()) {
                    DistributedStore.Options storeOptions = new DistributedStore.Options()
                            .type(DistributedStore.StoreType.FILE)
                            .masterPort(30800 + r);

                    DistributedStore store = DistributedStore.create(storeOptions, r, config.worldSize);

                    ProcessGroupWrapper.Options pgOptions = new ProcessGroupWrapper.Options()
                            .backend(useGpu ? ProcessGroupWrapper.BackendType.NCCL : ProcessGroupWrapper.BackendType.GLOO);

                    ProcessGroupWrapper pg = ProcessGroupWrapper.create(pgOptions, r, config.worldSize, store);

                    BenchmarkNet model = createModelOnDevice(config.inputSize, config.hiddenSize, config.numClasses, useGpu, r);

                    FSDPTrainer fsdp = FSDPTrainer.builder()
                            .module(model)
                            .processGroup(pg)
                            .build();

                    Optimizer optimizer = new Adam(model.parameters(), new AdamOptions(1e-3));

                    for (int step = 0; step < config.testSteps; step++) {
                        Tensor input = torch.randn(new long[]{config.batchSize, config.inputSize});
                        Tensor target = torch.randint(config.numClasses, new long[]{config.batchSize});

                        if (useGpu) {
                            input = input.to(new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Float);
                            target = target.to(new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Long);
                        }

                        Tensor output = fsdp.forward(input);
                        Tensor loss = torch.cross_entropy(output, target);

                        optimizer.zero_grad();
                        loss.backward();
                        optimizer.step();

                        // Check gradient validity safely
                        try {
                            TensorVector params = model.parameters();
                            var it = params.begin();
                            var end = params.end();
                            while(!it.equals(end)){
                                Tensor p = it.get();
                                if (p != null && !p.isNull()) {
                                    try {
                                        Tensor grad = p.grad();
                                        if (grad != null && !grad.isNull() && grad.defined()) {
                                            totalGradients.incrementAndGet();
                                            float norm = grad.norm().item_float();
                                            if (!Float.isNaN(norm) && !Float.isInfinite(norm) && norm < 100) {
                                                validGradients.incrementAndGet();
                                            }
                                        }
                                    } catch (Exception e) { /* skip */ }
                                }
                                it.increment();
                            }
                        } catch (Exception e) { /* skip */ }
//                        for (Tensor p : model.parameters()) {
//                            if (p.grad().defined()) {
//                                totalGradients.incrementAndGet();
//                                float norm = p.grad().norm().item_float();
//                                if (!Float.isNaN(norm) && !Float.isInfinite(norm) && norm < 100) {
//                                    validGradients.incrementAndGet();
//                                }
//                            }
//                        }

                        loss.close();
                        input.close();
                        target.close();
                    }

                    if (validGradients.get() == totalGradients.get() && totalGradients.get() > 0) {
                        successCount.incrementAndGet();
                    }

                    fsdp.close();
                    pg.close();
                    store.close();
                } catch (Exception e) {
                    System.err.println("Error in gradient correctness test: " + e.getMessage());
                }
            });
        }

        waitForCompletion();

        long duration = System.currentTimeMillis() - start;
        double validity = totalGradients.get() > 0
                ? (double) validGradients.get() / totalGradients.get() * 100
                : 0;
        boolean passed = successCount.get() == config.worldSize && validity > 95;

        return new TestResult(name, passed,
                String.format("valid gradients: %d/%d (%.1f%%)",
                        validGradients.get(), totalGradients.get(), validity),
                validity, duration);
    }

    static TestResult testShardingStrategies(String name, TestConfig config, boolean useGpu) {
        long start = System.currentTimeMillis();
        Map<FSDPTrainer.ShardingStrategy, Boolean> results = new ConcurrentHashMap<>();

        for (FSDPTrainer.ShardingStrategy strategy : FSDPTrainer.ShardingStrategy.values()) {
            final FSDPTrainer.ShardingStrategy s = strategy;
            AtomicBoolean success = new AtomicBoolean(false);

            for (int rank = 0; rank < config.worldSize; rank++) {
                final int r = rank;
                CompletableFuture.runAsync(() -> {
                    try (PointerScope scope = new PointerScope()) {
                        DistributedStore.Options storeOptions = new DistributedStore.Options()
                                .type(DistributedStore.StoreType.FILE)
                                .masterPort(30900 + r + s.ordinal() * 100);

                        DistributedStore store = DistributedStore.create(storeOptions, r, config.worldSize);

                        ProcessGroupWrapper.Options pgOptions = new ProcessGroupWrapper.Options()
                                .backend(useGpu ? ProcessGroupWrapper.BackendType.NCCL : ProcessGroupWrapper.BackendType.GLOO);

                        ProcessGroupWrapper pg = ProcessGroupWrapper.create(pgOptions, r, config.worldSize, store);

                        BenchmarkNet model = createModelOnDevice(config.inputSize, config.hiddenSize, config.numClasses, useGpu, r);

                        FSDPTrainer fsdp = FSDPTrainer.builder()
                                .module(model)
                                .processGroup(pg)
                                .shardingStrategy(s)
                                .build();

                        Optimizer optimizer = new Adam(model.parameters(), new AdamOptions(1e-3));

                        // Run a few steps
                        for (int step = 0; step < 5; step++) {
                            Tensor input = torch.randn(new long[]{config.batchSize, config.inputSize});
                            Tensor target = torch.randint(config.numClasses, new long[]{config.batchSize});

                            if (useGpu) {
                                input = input.to(new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Float);
                                target = target.to(new Device(torch.DeviceType.CUDA, (byte) r), torch.ScalarType.Long);
                            }

                            Tensor output = fsdp.forward(input);
                            Tensor loss = torch.cross_entropy(output, target);

                            optimizer.zero_grad();
                            loss.backward();
                            optimizer.step();

                            loss.close();
                            input.close();
                            target.close();
                        }

                        success.set(true);

                        fsdp.close();
                        pg.close();
                        store.close();
                    } catch (Exception e) {
                        System.err.println("Error testing strategy " + s + ": " + e.getMessage());
                    }
                });
            }

            waitForCompletion();

            // Wait for all ranks
            try { Thread.sleep(200); } catch (InterruptedException ignored) {}

            results.put(s, success.get());
        }

        long duration = System.currentTimeMillis() - start;
        boolean allPassed = results.values().stream().allMatch(Boolean::booleanValue);

        String strategies = results.entrySet().stream()
                .map(e -> e.getKey() + ":" + (e.getValue() ? "OK" : "FAIL"))
                .collect(Collectors.joining(", "));

        return new TestResult(name, allPassed,
                strategies,
                allPassed ? 1 : 0, duration);
    }

    static void waitForCompletion() {
        try {
            Thread.sleep(2000);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    static void printResults(List<TestResult> results) {
        System.out.println();
        System.out.println("=".repeat(90));
        System.out.println("FSDP Integration Test Results");
        System.out.println("=".repeat(90));
        for (TestResult r : results) {
            System.out.println(r);
        }
    }

    static BenchmarkNet createModelOnDevice(int inputSize, int hiddenSize, int numClasses, boolean useGpu, int rank) {
        BenchmarkNet model = new BenchmarkNet(inputSize, hiddenSize, numClasses);
        if (useGpu) {
            Device device = new Device(torch.DeviceType.CUDA, (byte) rank);
            model.to(device, true);
            // Collect params first, then move to GPU
            List<Tensor> paramList = new ArrayList<>();
            try {
                TensorVector params = model.parameters();
                var it = params.begin();
                var end = params.end();
                while(!it.equals(end)){
                    Tensor p = it.get();
                    if (p != null && !p.isNull()) paramList.add(p);
                    it.increment();
                }
            } catch (Exception e) { /* ignore */ }
            // Now move each to GPU
            for (Tensor p : paramList) {
                try {
                    if (p.device().type() != torch.DeviceType.CUDA) {
                        Tensor gpuCopy = torch.empty(p.shape()).to(device, torch.ScalarType.Float);
                        gpuCopy.copy_(p);
                        p.set_(gpuCopy);
                    }
                } catch (Exception e) { /* ignore */ }
            }
            // Move the Sequential's internal modules too
            try {
                SequentialImpl seq = model.asSequential();
                if (seq != null) {
                    seq.to(device, true);
                }
            } catch (Exception e) { /* ignore */ }
        }
        return model;
    }
}

package torch;

import org.bytedeco.pytorch.*;
import org.bytedeco.pytorch.global.torch;
import org.bytedeco.pytorch.SequentialImpl;

/**
 * BenchmarkNet - Simple MLP for benchmarking.
 * Extends SequentialImpl for JavaCPP compatibility with distributed training.
 */

import org.bytedeco.pytorch.*;
import org.bytedeco.pytorch.Module;

/**
 * BenchmarkNet - Simple MLP for benchmarking. Must be in torch package for javacpp compatibility.
 */
public class BenchmarkNet extends Module{
    private final SequentialImpl layers;

    public BenchmarkNet(long inputSize, long hiddenSize, long outputSize) {
        layers = register_module("layers", new SequentialImpl());
        layers.push_back(new LinearImpl(inputSize, hiddenSize));
        layers.push_back(new ReLUImpl());
        layers.push_back(new LinearImpl(hiddenSize, hiddenSize));
        layers.push_back(new ReLUImpl());
        layers.push_back(new LinearImpl(hiddenSize, outputSize));
    }

    public Tensor forward(Tensor x) {
        return layers.forward(x);
    }
}
//public class BenchmarkNet extends SequentialImpl {
//
//    public BenchmarkNet(long inputSize, long hiddenSize, long outputSize) {
//        push_back(new LinearImpl(inputSize, hiddenSize));
//        push_back(new ReLUImpl());
//        push_back(new LinearImpl(hiddenSize, hiddenSize));
//        push_back(new ReLUImpl());
//        push_back(new LinearImpl(hiddenSize, outputSize));
//    }
//}

[mullerhai]

Lack of overloaded forward methods in the Module class will become a critical system flaw.
I think ， @saudet