import torch
from tiny_fsdp.core import SGD, AdamW
from example.model import GPT2Model, GPTConfigs

model = GPT2Model(GPTConfigs.gpt2)
optimizer = AdamW(model.named_parameters(), lr=1e-4)

# Standard training loop
for batch in dataloader:
    output = model(batch)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

import torch.distributed as dist
from tiny_fsdp.core import DDP, DDPAdamW

# Initialize distributed training
dist.init_process_group(backend='nccl')

# Wrap model with DDP
model = GPT2Model(config)
model = DDP(model)

# Use DDP-aware optimizer
optimizer = DDPAdamW(model.named_parameters(), lr=1e-4)

# Training loop
model.require_backward_grad_sync = True  # Enable gradient sync
for batch in dataloader:
    loss = model(batch)
    loss.backward()  # Gradients automatically all-reduced
    optimizer.step()

from tiny_fsdp.core import Zero3, Zero3AdamW, zero3_partition_tensors

# Partition model parameters across ranks
with torch.device('meta'):
    model = GPT2Model(config)
    
parts, _ = zero3_partition_tensors(
    OrderedDict(model.named_parameters()),
    ranks_map=[f"cuda:{i}" for i in range(world_size)],
    evenness_priority=0
)

# Wrap with Zero3
model = GPT2Model(config)
model = Zero3(model, parts)

# Use Zero3-aware optimizer
optimizer = Zero3AdamW(
    model.module.named_parameters(), 
    lr=1e-4, 
    param_part_table=parts,
    ranks_map=[f"cuda:{i}" for i in range(world_size)]
)

from tiny_fsdp.core import FSDP, FSDPAdamW

# Wrap model with FSDP (automatic parameter sharding)
model = GPT2Model(config)
model = FSDP(model, world_size=world_size, rank=rank)

# Use FSDP-aware optimizer (works with sharded parameters)
optimizer = FSDPAdamW(model.named_parameters(), lr=1e-4)

# Training loop
for batch in dataloader:
    output = model(batch)  # Parameters auto-gathered
    loss = criterion(output, target)
    loss.backward()        # Gradients auto-scattered
    optimizer.step()       # Update local shards

# 1. Implement the module wrapper
class MyCustomModule(base_module.MyCustomModule):
    def forward_callback(self, ctx, *args):
        # Custom forward logic with parameter sync
        pass
        
    def backward_callback(self, ctx, grad_output):
        # Custom backward logic with gradient handling
        pass

# 2. Register in the strategy wrapper
_supported_modules = {
    nn.MyCustomModule: MyCustomModule,
    # ... other modules
}

# Enable gradient checkpointing
model.gradient_checkpointing = True

# Tune communication overlap
model.enable_async_communication = True

# Configure precision
model.half()  # Use FP16 for memory efficiency