memory_optimization/README.md

Memory Optimization Module

Unified memory management system with shared Qwen model integration for zero memory overhead.

Module Structure

memory_optimization/
├── __init__.py              # Module exports and convenience functions
├── config.py                # MemoryOptimizationConfig
├── manager.py               # UnifiedMemoryManager
├── tensor_pool.py           # TensorPool
├── model_cache.py           # ModelCache (uses shared Qwen model)
├── cleanup.py               # MemoryCleanup
└── README.md                # This file

Features

✅ Shared Model Integration

• ModelCache: Uses shared Qwen model for zero memory overhead
• Automatic fallback to cached models if shared model unavailable
• Prevents model duplication across modules

✅ CUDA Optimization

• All operations run on CUDA when available
• Efficient tensor pooling
• Adaptive memory cleanup

✅ Self-Contained Modules

• Each component is independent
• Easy to test and benchmark
• Clean separation of concerns

Usage

Basic Usage

from memory_optimization import (
    UnifiedMemoryManager,
    MemoryOptimizationConfig,
    get_unified_memory_manager
)

# Initialize with shared model
config = MemoryOptimizationConfig(
    use_shared_model=True,
    device="cuda"
)
manager = UnifiedMemoryManager(config)

# Get shared model (uses shared Qwen if available)
model = manager.get_shared_model("Qwen/Qwen3-0.6B", "transformer")

# Get optimized tensor
tensor = manager.get_tensor((10, 1024), dtype=torch.float32)

# Return tensor to pool
manager.return_tensor(tensor)

Convenience Functions

from memory_optimization import (
    get_shared_model,
    get_tensor,
    return_tensor,
    clear_memory,
    get_memory_stats
)

# Get shared model
model = get_shared_model("Qwen/Qwen3-0.6B", "transformer")

# Get tensor
tensor = get_tensor((10, 1024))

# Return tensor
return_tensor(tensor)

# Get stats
stats = get_memory_stats()

# Clear memory
clear_memory()

Integration with Shared Model

The module automatically detects and uses the shared Qwen model:

1. ModelCache: Uses shared Qwen model for transformers (zero memory overhead)
2. ModelCache: Uses shared Qwen tokenizer (zero memory overhead)
3. Automatic Fallback: Falls back to cached models if shared model unavailable

CUDA Compatibility

All components are CUDA-compatible:

• Automatic device detection
• Efficient GPU memory management
• Adaptive cleanup based on memory pressure

Configuration

See config.py for all configuration options. Key settings:

• use_shared_model: Enable shared Qwen model (default: True)
• device: Device to use ("cuda" or "cpu")
• memory_threshold: Memory usage threshold for cleanup
• max_pool_size: Maximum tensor pool size
• use_4bit_quantization: Enable 4-bit quantization

Memory Management

The system provides:

• Tensor Pooling: Reuse tensors to reduce allocations
• Model Caching: Share model instances across modules
• Adaptive Cleanup: Automatic memory cleanup based on pressure
• Emergency Cleanup: Force cleanup when memory is critical

Dependencies

• PyTorch (CUDA support recommended)
• Transformers
• BitsAndBytes (for 4-bit quantization, optional)