import torch
from ramtorch import Linear

# Standard PyTorch approach (high GPU memory usage)
# linear = torch.nn.Linear(1000, 1000)

# RamTorch approach (low GPU memory usage)
linear = Linear(1000, 1000, device="cuda")

# Use exactly like a normal PyTorch layer
x = torch.randn(32, 1000, device="cuda")
output = linear(x)  # Parameters automatically transferred from CPU to GPU

import torch.nn as nn
from ramtorch import Linear

class MemoryEfficientModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            Linear(1000, 2000),
            nn.ReLU(),
            Linear(2000, 2000),
            nn.ReLU(),
            Linear(2000, 100)
        )
    
    def forward(self, x):
        return self.layers(x)

model = MemoryEfficientModel()

from ramtorch.helpers import replace_linear_with_ramtorch

# Your existing PyTorch model
model = YourExistingModel()

# Replace all nn.Linear layers with RamTorch Linear layers
model = replace_linear_with_ramtorch(model, rank=0)
model = model.to("cuda:0")

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.utils.data import DataLoader, DistributedSampler

from ramtorch import AdamW
from ramtorch.helpers import replace_linear_with_ramtorch
from ramtorch.zero1 import create_zero_param_groups, broadcast_zero_params
from ramtorch.zero2 import setup_grad_sharding_hooks

def train(rank, world_size, model):
    # Setup distributed process group
    dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    
    # Replace Linear layers with RamTorch (shares CPU memory across workers)
    model = replace_linear_with_ramtorch(model, rank)
    model.to(rank)
    
    # Setup ZeRO-1: Shard optimizer states across workers
    # Each worker only maintains optimizer states for a subset of parameters
    all_params = list(model.parameters())
    param_groups = [{'params': all_params, 'lr': 1e-3, 'weight_decay': 0.01}]
    rank_param_groups = create_zero_param_groups(param_groups, world_size)
    
    # Setup ZeRO-2: Shard gradients across workers
    # Gradients are partitioned and only linked on the worker responsible for them
    setup_grad_sharding_hooks(rank_param_groups, rank)
    
    # Each worker's optimizer only handles its shard
    optimizer = AdamW(rank_param_groups[rank])
    
    # Scheduler works normally
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
    
    # Setup distributed data loading
    dataset = YourDataset()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = DataLoader(dataset, batch_size=32, sampler=sampler)
    
    # Training loop
    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # Important for proper shuffling
        
        for batch in loader:
            inputs, targets = batch
            inputs = inputs.to(rank)
            targets = targets.to(rank)
            
            # Forward and backward pass
            outputs = model(inputs)
            loss = criterion(outputs, targets) / world_size  # Scale loss
            
            # Synchronize before backward to ensure all workers are ready
            torch.cuda.synchronize()
            loss.backward()
            torch.cuda.synchronize()
            
            # Update parameters (each worker updates only its shard)
            optimizer.step()
            
            # Broadcast updated parameters to all workers
            # RamTorch parameters are already shared via CPU memory,
            # but standard PyTorch parameters need explicit broadcasting
            broadcast_zero_params(rank_param_groups)
            
            scheduler.step()
            
            # IMPORTANT: Use model.zero_grad(), not optimizer.zero_grad()
            # Each worker handles partial gradients, so we need to zero
            # gradients at the model level to properly clear all workers' buffers
            model.zero_grad()
            
            torch.cuda.synchronize()
    
    dist.destroy_process_group()

if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    
    # Model must be instantiated BEFORE spawning GPU workers
    # RamTorch shares CPU tensors across workers, so the model needs to exist
    # in the parent process before forking to enable proper memory sharing
    model = YourModel()
    
    mp.spawn(train, args=(world_size, model), nprocs=world_size)