xmem-mtier 0.1.2

High-level Python wrapper for MTier memory management with PyTorch integration

xmem

High-level Python wrapper for MTier memory management with PyTorch integration.

Overview

xmem enables seamless offloading of PyTorch tensors to MTier remote memory tiers, allowing you to work with datasets and models that exceed GPU memory capacity.

Key Features:

• Simple API: xmem.offload() and .prefetch()
• Seamless PyTorch integration with CUDA support
• Async operations for overlapping computation and I/O
• Automatic memory management

Quick Installation

# Install prerequisites
pip install torch torchvision cuda-python>=12.0.0

# Install xmem (development mode)
cd /path/to/memory-tier-emu-framework/python
pip install -e .

Requirements:

• Python >= 3.8
• PyTorch >= 2.0.0 with CUDA support
• NVIDIA Hopper+ GPU (H100, H200, etc.) with fabric memory support
• libmtier.so shared library (set MTIER_LIB_PATH if not in default location)

Quick Start

import torch
import xmem

# Create a tensor on GPU
tensor = torch.randn(1000, 1000, device='cuda')

# Offload to remote memory
remote_tensor = xmem.offload(tensor, bank_id=1)

# Prefetch back to GPU
reconstructed = remote_tensor.prefetch()

# Verify data integrity
assert torch.equal(tensor, reconstructed)

Async Operations

# Non-blocking offload
remote_tensor = xmem.offload(tensor, bank_id=1, non_blocking=True)

# Do other work while transfer happens...

# Non-blocking prefetch
remote_tensor.prefetch(non_blocking=True)

# Do other work...

# Get the result
reconstructed = remote_tensor.wait()

Documentation

📚 Full Python API Documentation

The comprehensive documentation includes:

• Detailed installation guide with troubleshooting
• Complete API reference for all functions and classes
• Advanced usage patterns and best practices
• Performance optimization tips
• Error handling and exception reference

Examples

See working examples in the examples/ directory:

• pytorch_usage.py - Basic synchronous operations
• pytorch_async_usage.py - Async operations with pipelining

cd examples
python pytorch_usage.py
python pytorch_async_usage.py

API Summary

Core Functions

• xmem.offload(tensor, bank_id, non_blocking=False, stream=None) → XMemTensorMTier
• xmem.is_available() → Check if MTier is available
• xmem.mbank_count() → Get number of memory banks
• xmem.memory_available(bank_id) → Get available memory in bytes
• xmem.memory_allocated(bank_id) → Get allocated memory in bytes

XMemTensorMTier Methods

• .prefetch(dst_tensor=None, device=None, reuse=False, non_blocking=False, stream=None) → Retrieve tensor from remote memory
• .wait(timeout=None) → Wait for async operation to complete
• .is_ready() → Check if async operation is complete
• Properties: shape, dtype, size, remote_addr, device

Exceptions

All exceptions inherit from MTierError:

• MTierNotInitializedError - SDK not initialized
• MTierOutOfMemoryError - Remote memory allocation failed
• MTierInvalidArgumentError - Invalid argument
• MTierServiceUnavailableError - MTier service unavailable

See full documentation for complete exception hierarchy.

Development

# Install in development mode
pip install -e ".[dev]"

# Build wheel
pip install build
python -m build

# Run tests
pytest tests/

Support

For detailed documentation, examples, and best practices, see the Python API Documentation.

For issues or questions, please contact the development team.

---

License: Proprietary - All Rights Reserved