Here is a reproducible example with the second Error (inplace operation), assuming you have two GPUs available (what’s interesting, now it doesn’t throw the first error related to the ‘read only once’, even though the gradient checkpointing is still turned on). Some parts of the code were simplified, just to provide a working example (i.e. random data creation and sampling through it - I know it could’ve been done more carefully, but wanted to provide working example as quickly as possible):
import os

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.multiprocessing as mp
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

from transformers import AutoConfig, AutoModel

def run_training(rank):
    torch.autograd.set_detect_anomaly(True)
    torch.distributed.init_process_group(backend="nccl", rank=rank)
    
    device = torch.device('cuda', rank)
    
    length_inp = 256
    data = [{'sent1': {'input_ids': torch.randint(0, 200, (16, length_inp)), 'attention_mask': torch.ones(16, length_inp)}, 
             'sent2': {'input_ids': torch.randint(0, 200, (16, length_inp)), 'attention_mask': torch.ones(16, length_inp)},
             'sent3': {'input_ids': torch.randint(0, 200, (16, length_inp)), 'attention_mask': torch.ones(16, length_inp)}
            }]*8
    
    train_sampler = DistributedSampler(
                    data,
                    num_replicas=torch.cuda.device_count(),
                    rank=rank,
                    seed=44,
                )
    
    dll = DataLoader(
            data,
            batch_size=1,
            sampler=train_sampler,
            drop_last=False,
            num_workers=2,
            pin_memory=True,
        )

    config = AutoConfig.from_pretrained('distilroberta-base')
    model = AutoModel.from_pretrained('distilroberta-base', config=config, add_pooling_layer=False)
    model.to(device)
    model.gradient_checkpointing_enable()
    
    model = nn.parallel.DistributedDataParallel(model, device_ids=[rank], output_device=rank, find_unused_parameters=False)
    optimizer = torch.optim.Adam(model.parameters())
    metric = lambda x,y: 1.0 - F.cosine_similarity(x, y)
    criterion = nn.TripletMarginWithDistanceLoss(distance_function=metric, margin=0.2, reduction='none')
    
    for n, b in enumerate(dll):
        print(n)
        model.zero_grad()
        sent1 = {k:v.squeeze(0) for k,v in b['sent1'].items()}
        sent2 = {k:v.squeeze(0) for k,v in b['sent2'].items()}
        sent3 = {k:v.squeeze(0) for k,v in b['sent3'].items()}
        
        emb1 = model(**sent1)[0][:, 0, :]
        emb2 = model(**sent2)[0][:, 0, :]
        emb3 = model(**sent3)[0][:, 0, :]
        
        losses = criterion(emb1, emb2, emb3)
        loss = losses.mean()
        loss.backward()
        
        optimizer.step()
        print('Model device: {}, loss device: {}, loss: {}'.format(model.device, loss.device, loss))


def main():
    world_size = torch.cuda.device_count()
    os.environ["MASTER_PORT"] = '1234'
    os.environ["MASTER_ADDR"] = '127.0.0.1'
    os.environ["WORLD_SIZE"] = str(world_size)
    mp.spawn(run_training,
        nprocs=world_size,
        join=True)

if __name__ == "__main__":
    main()