First of all, thank you for the great work on javacpp-pytorch.
We have successfully implemented FSDP (Fully Sharded Data Parallel) distributed training using javacpp-pytorch, but we are facing a critical off-heap memory leak issue.
Problem Description
Heap memory: stable at ~8MB (normal)
Off-heap (native/CUDA) memory: keeps growing with each training step, until OOM crashes
We are processing batches in loops, and the memory is never fully released
This happens in both single GPU and multi-GPU FSDP distributed training
What we have already tried
Using PointerScope explicitly for all local Tensors
Calling tensor.close() manually after use
Calling tensorVector.close() after use
Trying to deallocate outputs and intermediate tensors
Using System.gc() periodically (not effective)
But off-heap memory still grows without bounds.
Our conclusion
It seems there are still some underlying C++ / native tensor resources not being properly managed or deallocated during training loops, especially under FSDP distribution.
Ask
Could you give us some advice or fix regarding:
How to fully clean up native memory in training loops
Proper usage of PointerScope in FSDP distributed training
Any missing cleanup methods for native tensors/modules/optimizers
This issue blocks us from running long-term training jobs.
Thank you very much for your help!
package torch;
import org.bytedeco.javacpp.PointerScope;
import org.bytedeco.javacpp.chrono.Milliseconds;
import org.bytedeco.pytorch.*;
import org.bytedeco.pytorch.Module;
import org.bytedeco.pytorch.global.torch;
import org.bytedeco.pytorch.global.torch_cuda;
import java.util.*;
public class FSDOTrainingCompleteV3 {
// 追踪张量创建（你自己的计数器，100% 可用）
private static long tensorCount = 0;
public static void incTensor() { tensorCount++; }
public static void decTensor() { tensorCount--; }
public static long getTensorCount() { return tensorCount; }
public static void logTensorCount(String tag) {
System.out.println("\n🔥 当前存活 Tensor 数量 = " + tensorCount + " | 位置: " + tag);
if (tensorCount > 5000) {
System.out.println("⚠️ ⚠️ ⚠️ 严重内存泄漏！张量太多！！！");
}
}
// -------------------------------------------------------------------------
// Main 可用
// -------------------------------------------------------------------------
public static void main(String[] args) {
int rank = 0;
int worldSize = 1;
String storePath = "/tmp/fsdp_store";
System.out.println("FSDP 训练启动 | rank=" + rank + " worldSize=" + worldSize);
// 设置随机种子，保证可复现
torch.manual_seed(42L);
// 模型
long inputSize = 784;
long hidden = 128;
long classes = 10;
final long vocabSize = 10000;
final long dModel = 512;
final long nhead = 8;
final long numLayers = 6;
torch.manual_seed(42L);
// Create model
TransformerModelV2 model = new TransformerModelV2(vocabSize, dModel, nhead, numLayers);
// model.to(device,false);
// SimpleNetComplete model = new SimpleNetComplete(inputSize, hidden, classes);
model.to(new Device(torch.DeviceType.CPU,(byte) rank),false);
model.parameters();
// FSDP
FSDPTrainerComplete trainer = new FSDPTrainerComplete(model, rank, worldSize, storePath);
// 🔥 优化器绑定分片参数（实际未使用，我们手动 SGD）
Optimizer optimizer = new SGD(new TensorVector(trainer.shardedParam), new SGDOptions(0.01));
// 采样器
DistributedSamplerCompleteV2 sampler = new DistributedSamplerCompleteV2(100, rank, worldSize, true, 42);
// ✅ 关键修复：使用归一化的随机数据（而不是 arange[0..783] 这种大数值）
// 用 randn 生成均值 0、方差 1 的数据，模拟真实场景（如归一化的 MNIST）
// List<Tensor> data = new ArrayList<>();
// List<Tensor> label = new ArrayList<>();
var numEpochs = 5;
model.train(true);
for (int epoch = 0; epoch < numEpochs; epoch++) {
Tensor x = torch.arange(new Scalar(0), new Scalar(784)).reshape(784).to(torch.ScalarType.Long);
// 标签保持基于x生成，不随机
Tensor y = torch.tensor(new long[]{(x.sum().item_long() % 10)}).to(torch.ScalarType.Long);
for (int i = 0; i < 100; i++) {
try(PointerScope scope = new PointerScope()) {
// ✅ 输入：标准正态分布，数值范围 ~[-3, 3]，避免数值爆炸
System.out.println("\n=== Epoch " + (epoch + 1) + "/" + numEpochs + " === "+ i);
sampler.setEpoch(epoch);
trainer.trainStep(x, y, optimizer);
// System.runFinalization();
}
}
x.close();
y.close();
System.gc();
}
System.out.println("[Rank " + rank + "] 数据初始化完成: 100 samples, input shape=[784], 10 classes");
// 训练
// trainer.train(optimizer, data, label, sampler, 5);
System.out.println("训练完成！");
}
public static class DistributedSamplerCompleteV2 {
private long numSamples;
private int rank;
private int worldSize;
private long numSamplesPerRank;
private List<Long> indices;
private boolean shuffle;
private int seed;
public DistributedSamplerCompleteV2(long totalSamples, int rank, int worldSize,
boolean shuffle, int seed) {
this.numSamples = totalSamples;
this.rank = rank;
this.worldSize = worldSize;
this.shuffle = shuffle;
this.seed = seed;
this.numSamplesPerRank = (numSamples + worldSize - 1) / worldSize;
this.indices = new ArrayList<>();
System.out.println("[Rank " + rank + "] Sampler initialized: " +
numSamplesPerRank + " samples per rank");
}
public void setEpoch(int epoch) {
indices.clear();
List<Long> allIndices = new ArrayList<>();
for (long i = 0; i < numSamples; i++) allIndices.add(i);
if (shuffle) {
Random rand = new Random(seed + epoch);
Collections.shuffle(allIndices, rand);
}
long totalSize = numSamplesPerRank * worldSize;
while (allIndices.size() < totalSize) {
allIndices.add(allIndices.get((int) (allIndices.size() % numSamples)));
}
for (long i = rank; i < allIndices.size(); i += worldSize) {
indices.add(allIndices.get((int) i));
}
}
public List<Long> getIndices() { return indices; }
public long size() { return numSamplesPerRank; }
}
public static class TransformerModel extends Module {
private final EmbeddingImpl embedding;
private final TransformerEncoderImpl encoder;
private final LinearImpl fc;
public TransformerModel(long vocabSize, long dModel, long nhead, long numLayers) {
this.embedding = register_module("embedding", new EmbeddingImpl(vocabSize, dModel));
var transformerOpt = new TransformerEncoderLayerOptions(dModel, nhead);
transformerOpt.dim_feedforward().put(dModel * 4);
transformerOpt.dropout().put(0.1);
TransformerEncoderLayerImpl encoderLayer = new TransformerEncoderLayerImpl(transformerOpt);
var encoderOpt =new TransformerEncoderOptions(transformerOpt, numLayers);
this.encoder = register_module("encoder", new TransformerEncoderImpl(encoderOpt));
this.fc = register_module("fc", new LinearImpl(dModel, vocabSize));
}
public Tensor forward(Tensor x) {
var dim =Math.sqrt(embedding.options().embedding_dim().get());
x = embedding.forward(x).mul(new Scalar(dim));
x = encoder.forward(x);
x = fc.forward(x);
return x;
}
}
public static class TransformerModelV3 extends Module {
private final EmbeddingImpl embedding;
private final TransformerEncoderImpl encoder;
private final LinearImpl fc;
public TransformerModelV3(long vocabSize, long dModel, long nhead, long numLayers) {
this.embedding = register_module("embedding", new EmbeddingImpl(vocabSize, dModel));
// ✅ 修复：标准 API 构造，避免未初始化参数
TransformerEncoderLayerOptions transformerOpt = new TransformerEncoderLayerOptions(dModel, nhead);
transformerOpt.dim_feedforward().put(dModel * 4);
transformerOpt.dropout().put(0.1);
transformerOpt.activation().put(new kReLU());
TransformerEncoderLayerImpl encoderLayer = new TransformerEncoderLayerImpl(transformerOpt);
var encoderOpt =new TransformerEncoderOptions(transformerOpt, numLayers);
this.encoder = register_module("encoder", new TransformerEncoderImpl(encoderOpt));
this.fc = register_module("fc", new LinearImpl(dModel, vocabSize));
}
public Tensor forward(Tensor x) {
double dim = Math.sqrt(embedding.options().embedding_dim().get());
x = embedding.forward(x).mul(new Scalar(dim));
x = encoder.forward(x);
x = fc.forward(x);
return x;
}
}
public static class TransformerModelV4 extends Module {
private final EmbeddingImpl embedding;
private final TransformerEncoderImpl encoder;
private final LinearImpl fc;
private final LinearImpl inputProjection; // 🔥 最小改动：加一层投影
public TransformerModelV4(long vocabSize, long dModel, long nhead, long numLayers) {
this.embedding = register_module("embedding", new EmbeddingImpl(vocabSize, dModel));
// 完全沿用你自己的写法！
TransformerEncoderLayerOptions transformerOpt = new TransformerEncoderLayerOptions(dModel, nhead);
transformerOpt.dim_feedforward().put(dModel * 4);
transformerOpt.dropout().put(0.1);
transformerOpt.activation().put(new kReLU());
TransformerEncoderLayerImpl encoderLayer = new TransformerEncoderLayerImpl(transformerOpt);
var encoderOpt = new TransformerEncoderOptions(transformerOpt, numLayers);
this.encoder = register_module("encoder", new TransformerEncoderImpl(encoderOpt));
this.fc = register_module("fc", new LinearImpl(dModel, vocabSize));
// 🔥 最小修复：把 784 维输入映射到 dModel，解决断言错误
this.inputProjection = register_module("inputProjection", new LinearImpl(784, dModel));
}
public Tensor forward(Tensor x) {
// 🔥 核心修复：输入 x [784] → 先投影到 dModel，再进 Transformer
x = inputProjection.forward(x); // [784] → [dModel]
// 🔥 再升维成 Transformer 要求格式 [seq_len=1, batch=1, d_model]
x = x.unsqueeze(0).unsqueeze(0);
// 你原有逻辑完全不变
double dim = Math.sqrt(embedding.options().embedding_dim().get());
x = embedding.forward(x).mul(new Scalar(dim));
x = encoder.forward(x);
x = x.squeeze(0).squeeze(0); // 降维回去
x = fc.forward(x);
return x;
}
}
public static class TransformerModelV2 extends Module {
private final EmbeddingImpl embedding;
private final TransformerEncoderImpl encoder;
private final LinearImpl fc;
public TransformerModelV2(long vocabSize, long dModel, long nhead, long numLayers) {
this.embedding = register_module("embedding", new EmbeddingImpl(vocabSize, dModel));
// 完全用你原来的 API，不动！
TransformerEncoderLayerOptions transformerOpt = new TransformerEncoderLayerOptions(dModel, nhead);
transformerOpt.dim_feedforward().put(dModel * 4);
transformerOpt.dropout().put(0.1);
transformerOpt.activation().put(new kReLU());
TransformerEncoderLayerImpl encoderLayer = new TransformerEncoderLayerImpl(transformerOpt);
var encoderOpt = new TransformerEncoderOptions(transformerOpt, numLayers);
this.encoder = register_module("encoder", new TransformerEncoderImpl(encoderOpt));
this.fc = register_module("fc", new LinearImpl(dModel, vocabSize));
}
// ✅ 终极修复：维度 + 类型 + batch 完全匹配 cross_entropy
public Tensor forward(Tensor x) {
// 输入 x: [784] float → 先转成 token index (long)
x = x.to(torch.ScalarType.Long);
// embedding 输出: [784, d_model]
double dim = Math.sqrt(embedding.options().embedding_dim().get());
x = embedding.forward(x).mul(new Scalar(dim));
// Transformer 要求: [seq_len, batch, d_model]
x = x.unsqueeze(1); // [784, 1, d_model]
// 过 encoder
x = encoder.forward(x);
// 输出变成 [1, vocab_size] 匹配 cross_entropy 要求
x = x.mean(0); // 全局池化 → [1, d_model]
x = fc.forward(x); // [1, vocab_size]
return x;
}
}
// -------------------------------------------------------------------------
// Model
// -------------------------------------------------------------------------
static class SimpleNetComplete extends Module {
private LinearImpl fc1, fc2, fc3;
public SimpleNetComplete(long inputSize, long hiddenSize, long numClasses) {
this.fc1 = register_module("fc1", new LinearImpl(inputSize, hiddenSize));
this.fc2 = register_module("fc2", new LinearImpl(hiddenSize, hiddenSize));
this.fc3 = register_module("fc3", new LinearImpl(hiddenSize, numClasses));
}
public Tensor forward(Tensor x) {
x = torch.relu(fc1.forward(x));
x = torch.relu(fc2.forward(x));
x = fc3.forward(x);
return x;
}
}
// -------------------------------------------------------------------------
// FSDP Trainer (FIXED VERSION)
// -------------------------------------------------------------------------
static class FSDPTrainerComplete {
// private final SimpleNetComplete model;
private final TransformerModelV2 model;
private final ProcessGroupGloo processGroup;
private final int rank;
private final int worldSize;
// 🔥 核心修复：FSDP 分片参数 + 梯度
private Tensor shardedParam;
private Tensor shardedGrad;
private Tensor shardedParamForOpt; // 用于优化器追踪的参数副本
// 模型参数结构（用于还原形状）
private List<Long> paramShapes;
private List<Long> paramNumels;
private long totalParamNumel;
public FSDPTrainerComplete(TransformerModelV2 model, int rank, int worldSize, String masterAddr) {
this.model = model;
this.rank = rank;
this.worldSize = worldSize;
// 1. 初始化进程组
FileStore store = new FileStore(masterAddr, worldSize);
ProcessGroupGloo.Options options = new ProcessGroupGloo.Options();
options.timeout(new Milliseconds(5000));
options.devices().push_back(ProcessGroupGloo.createDeviceForHostname("127.0.0.1"));
this.processGroup = new ProcessGroupGloo(store, rank, worldSize, options);
// 2. 收集模型参数结构信息
collectParamMetadata();
// 3. 广播初始参数
broadcastFullParameters();
// 4. 🔥 分片参数
shardParameters();
System.out.println("[Rank " + rank + "] FSDP 初始化完成，分片参数大小: " + shardedParam.numel());
}
// ---------------------------------------------------------------------
// 收集参数形状（用于后续 reshape 还原）
// ---------------------------------------------------------------------
private void collectParamMetadata() {
paramShapes = new ArrayList<>();
paramNumels = new ArrayList<>();
totalParamNumel = 0;
try (PointerScope scope = new PointerScope()) { // 🔥 自动释放
for (Tensor p : getModelParams(model)) {
paramShapes.addAll(Arrays.asList(p.sizes().get(0)));
paramNumels.add(p.numel());
totalParamNumel += p.numel();
}
}
}
// ---------------------------------------------------------------------
// 广播完整参数（所有 rank 保持一致）
// ---------------------------------------------------------------------
private void broadcastFullParameters() {
try (PointerScope scope = new PointerScope()) {
for (Tensor p : getModelParams(model)) {
// TensorVector vec = new TensorVector(p);
BroadcastOptions opts = new BroadcastOptions();
opts.rootRank(0);
processGroup.broadcast( new TensorVector(p), opts)._wait();
// vec.close();
}
}
if (rank == 0) System.out.println("[Rank 0] 参数广播完成");
}
// ---------------------------------------------------------------------
// 🔥 核心：参数分片
// =====================================================================
private void shardParameters() {
// try (PointerScope scope = new PointerScope()) {
// Flatten all parameters
Tensor flat = flattenParams(model);
// Calculate shard range
long shardSize = (totalParamNumel + worldSize - 1) / worldSize;
long start = rank * shardSize;
long end = Math.min(start + shardSize, totalParamNumel);
// 🔥 关键修复：创建可追踪梯度的参数（用于优化器）
shardedParamForOpt = flat.slice(0, new LongOptional(start), new LongOptional(end), 1).clone().detach();
shardedParamForOpt.requires_grad_(true); // 启用梯度追踪
// 保存数据引用（用于还原到模型）
shardedParam = shardedParamForOpt;
// 初始化梯度张量
shardedGrad = torch.zeros_like(shardedParamForOpt);
shardedGrad.requires_grad_(false);
// flat.close();
System.out.println("[Rank " + rank + "] 🔥 Sharded param created with grad tracking, size: " + shardedParamForOpt.numel());
// }
}
// ---------------------------------------------------------------------
// 🔥 核心：AllGather 还原完整参数
// ---------------------------------------------------------------------
private void allGatherToModel() {
int worldSize = this.worldSize;
long shardSize = shardedParam.numel();
// 1. 准备接收所有分片
List<Tensor> gathered = new ArrayList<>();
for (int i = 0; i < worldSize; i++) gathered.add(torch.empty(shardSize));
// 2. AllGather
// TensorVector out = new TensorVector(gathered.toArray(new Tensor[0]));
// TensorVector in = new TensorVector(shardedParam);
// processGroup.allgather(out, in, new AllgatherOptions())._wait();
processGroup.allgather( new TensorVector(gathered.toArray(new Tensor[0])), new TensorVector(shardedParam), new AllgatherOptions())._wait();
// 3. 拼接并还原到模型
Tensor full = torch.cat(new TensorVector(gathered.toArray(new Tensor[0])));
full = full.slice(0, new LongOptional(0), new LongOptional(totalParamNumel),1);
unflattenParams(full, model);
// out.close();
// in.close();
full.close();
gathered = new ArrayList<>();
}
// ---------------------------------------------------------------------
// 🔥 核心：反向传播后汇总梯度并写回分片
// ---------------------------------------------------------------------
private void reduceScatterGradients() {
// 1. 展平模型梯度
Tensor gradFlat = flattenGrads(model);
if (gradFlat == null) {
System.err.println("[Rank " + rank + "] Warning: Flattened gradients are null.");
return;
}
long shardSize = shardedParam.numel();
int worldSize = this.worldSize;
// 2. 切分梯度
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));
}
// 3. ReduceScatter 求和
Tensor localGrad = torch.empty_like(shardedParam);
// TensorVector outVec = new TensorVector(localGrad);
// TensorVector inVec = new TensorVector(splits.toArray(new Tensor[0]));
ReduceScatterOptions opts = new ReduceScatterOptions();
opts.reduceOp(new ReduceOp(ReduceOp.RedOpType.SUM));
processGroup.reduce_scatter(new TensorVector(localGrad), new TensorVector(splits.toArray(new Tensor[0])), opts)._wait();
// processGroup.reduce_scatter(outVec, inVec, opts)._wait();
// 4. 平均并保存到分片梯度
localGrad.div_(new Scalar(worldSize));
shardedGrad.data().copy_(localGrad);
// Debugging: Print gradient values
System.out.println("[Rank " + rank + "] Reduced gradient: " + localGrad);
// 5. 清空模型梯度（节省显存）
for (Tensor p : getModelParams(model)) {
if (p.grad().defined()) p.grad().zero_();
}
System.out.println("[Rank " + rank + "] Gradients reduced and scattered.");
gradFlat.close();
// outVec.close();
// inVec.close();
localGrad.close();
splits = new ArrayList<>();
}
// ---------------------------------------------------------------------
// 用分片参数+梯度执行一步优化
// ---------------------------------------------------------------------
// public void stepOptimizer(Optimizer optimizer) {
// // 把梯度赋值给分片参数
// shardedParam.grad().set_data(shardedGrad);
// // 优化器更新分片参数
// optimizer.step();
// // 清零梯度
// optimizer.zero_grad();
// }
public void stepOptimizer(Optimizer optimizer) {
// 🔥 关键修复：不要清零梯度，让它在下次迭代开始时重新初始化
double lr = 0.001; // ✅ 降低 learning rate 防止数值爆炸
double maxGradNorm = 1.0; // ✅ 梯度裁剪阈值
// 检查梯度是否有效
if (!shardedGrad.defined()) {
System.err.println("[Rank " + rank + "] Error: shardedGrad is not defined!");
return;
}
float gradNorm = shardedGrad.norm().item_float();
System.out.println("[Rank " + rank + "] 🔥 Applying gradient update with lr=" + lr);
System.out.println("[Rank " + rank + "] Gradient norm (before clip): " + gradNorm);
// ✅ 梯度裁剪：防止梯度爆炸导致 Dead ReLU
if (gradNorm > maxGradNorm) {
double scale = maxGradNorm / (gradNorm + 1e-6);
shardedGrad.mul_(new Scalar(scale));
System.out.println("[Rank " + rank + "] Gradient clipped, scale=" + scale + ", new norm=" + shardedGrad.norm().item_float());
}
// 🔥 手动参数更新：param = param - lr * grad
shardedParamForOpt.data().add_( shardedGrad.mul(new Scalar(-lr)) );
// ✅ 不在这里清零梯度，让它在下次 trainStep 开始时重新初始化
System.out.println("[Rank " + rank + "] Parameter updated successfully");
}
// ---------------------------------------------------------------------
// 训练一步（完全修复）
// =====================================================================
public void trainStep(Tensor input, Tensor target, Optimizer optimizer) {
// MemoryMonitor.logJvmMemory("trainStep 开始");
// logTensorCount("trainStep 开始");
try(PointerScope scope =new PointerScope()){
// ✅ 关键修复 1: 清零模型梯度（防止梯度累积）
for (Tensor p : getModelParams(model)) {
if (p.grad().defined()) p.grad().zero_();
}
// ✅ 关键修复 2: 重新初始化梯度容器
shardedGrad = torch.zeros_like(shardedParamForOpt);
shardedGrad.requires_grad_(false);
// ✅ 关键修复 3: 把分片参数写回模型（无论单机还是多机都需要）
writeShardedParamsToModel();
// 2. 前向
// Tensor output = model.forward(input);
// Tensor loss = torch.cross_entropy(output, target);
Tensor loss = torch.cross_entropy(model.forward(input), target);
System.out.println("[Rank " + rank + "] Loss before backward: " + loss.item_float());
// 3. 反向
loss.backward();
// 4. 检查梯度
Tensor modelGradFlat = flattenGrads(model);
if (modelGradFlat != null) {
System.out.println("[Rank " + rank + "] Model gradient norm: " + modelGradFlat.norm().item_float());
}
// 5. 汇总梯度
if (worldSize > 0) {
reduceScatterGradients();
} else {
// ✅ 单机情况：直接将模型梯度复制到 shardedGrad
Tensor gradFlat = flattenGrads(model);
if (gradFlat != null) {
shardedGrad.data().copy_(gradFlat);
System.out.println("[Rank " + rank + "] Sharded gradient copied, norm: " + shardedGrad.norm().item_float());
gradFlat.close();
} else {
System.err.println("[Rank " + rank + "] ERROR: Model gradients are null!");
return;
}
gradFlat.close();
}
// 6. 🔥 更新分片参数（手动SGD更新）
stepOptimizer(optimizer);
System.out.println("[Rank " + rank + "] Loss: " + loss.item_float());
// output.close();
loss.close();
// modelGradFlat.close();
}
// MemoryMonitor.logJvmMemory("trainStep 结束");
// logTensorCount("trainStep 结束");
// MemoryMonitor.forceClean(); // 强制释放
}
// ✅ 新增：把分片参数写回模型（FSDP 关键步骤）
// =====================================================================
private void writeShardedParamsToModel() {
if (worldSize > 0) {
// 多机：通过 AllGather 聚合分片
allGatherToModel();
} else {
// 单机：直接把分片参数写回模型
unflattenParams(shardedParamForOpt.data(), model);
}
}
// ---------------------------------------------------------------------
// 训练循环
// ---------------------------------------------------------------------
public void train(Optimizer optimizer,
List<Tensor> trainData,
List<Tensor> trainLabels,
DistributedSamplerCompleteV2 sampler,
int numEpochs) {
model.train(true);
for (int epoch = 0; epoch < numEpochs; epoch++) {
System.out.println("\n=== Epoch " + (epoch + 1) + "/" + numEpochs + " ===");
sampler.setEpoch(epoch);
for (long idx : sampler.getIndices()) {
int i = (int) (idx % trainData.size());
trainStep(trainData.get(i), trainLabels.get(i), optimizer);
}
}
}
// -------------------------------------------------------------------------
// 工具方法：展平参数 / 展平梯度 / 还原参数
// -------------------------------------------------------------------------
private static Tensor flattenParams(Module model) {
// try (PointerScope scope = new PointerScope()) {
List<Tensor> params = getModelParams(model);
List<Tensor> flat = new ArrayList<>();
for (Tensor p : params) flat.add(p.flatten());
return torch.cat(new TensorVector(flat.toArray(new Tensor[0])));
// }
}
private static Tensor flattenGrads(Module model) {
List<Tensor> grads = new ArrayList<>();
try (PointerScope scope = new PointerScope()) {
for (Tensor p : getModelParams(model)) {
if (!p.grad().defined()) return null;
grads.add(p.grad().flatten());
}
return torch.cat(new TensorVector(grads.toArray(new Tensor[0])));
}
}
private static void unflattenParams(Tensor flat, Module model) {
long offset = 0;
try (PointerScope scope = new PointerScope()) {
for (Tensor p : getModelParams(model)) {
long n = p.numel();
Tensor src = flat.slice(0, new LongOptional(offset), new LongOptional(offset + n),1).view(p.sizes());
p.data().copy_(src);
offset += n;
src.close();
p.close();
}
}
}
private static List<Tensor> getModelParams2(Module model) {
List<Tensor> res = new ArrayList<>();
try (PointerScope scope = new PointerScope()) { // 🔥 自动释放
TensorVector params = model.parameters();
for (long i = 0; i < params.size(); i++) {
res.add(params.get(i)); // 只拿引用，不新增内存
}
} // ✅ scope 退出，params 自动关闭
return res;
}
private static List<Tensor> getModelParams(Module model) {
List<Tensor> params = new ArrayList<>();
// try (PointerScope scope = new PointerScope()) { // 🔥 自动释放
var paramVector = model.parameters();
var begin = paramVector.begin();
var end = paramVector.end();
while (!begin.equals(end)) {
params.add(begin.get().data());
begin.increment();
}
// }
return params;
}
}
}
Hi JavaCPP team & @saudet
First of all, thank you for the great work on javacpp-pytorch.
We have successfully implemented FSDP (Fully Sharded Data Parallel) distributed training using javacpp-pytorch, but we are facing a critical off-heap memory leak issue.
Problem Description
Heap memory: stable at ~8MB (normal)
Off-heap (native/CUDA) memory: keeps growing with each training step, until OOM crashes
We are processing batches in loops, and the memory is never fully released
This happens in both single GPU and multi-GPU FSDP distributed training
What we have already tried
Using PointerScope explicitly for all local Tensors
Calling tensor.close() manually after use
Calling tensorVector.close() after use
Trying to deallocate outputs and intermediate tensors
Using System.gc() periodically (not effective)
But off-heap memory still grows without bounds.
Our conclusion
It seems there are still some underlying C++ / native tensor resources not being properly managed or deallocated during training loops, especially under FSDP distribution.
Ask
Could you give us some advice or fix regarding:
How to fully clean up native memory in training loops
Proper usage of PointerScope in FSDP distributed training
Any missing cleanup methods for native tensors/modules/optimizers
This issue blocks us from running long-term training jobs.
Thank you very much for your help!