Function test_offloading_integration

tests/test_custom_offloading_utils.py:341–375

()

# Integration test for offloading in a realistic scenario
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_offloading_integration():
    device = torch.device('cuda')
    # Create a mini model with 4 blocks
    model = SimpleModel(5)
    model.to(device)
    blocks = model.blocks
    
    # Initialize model offloader
    offloader = ModelOffloader(
        blocks=blocks,
        blocks_to_swap=2,
        device=device,
        debug=True
    )
    
    # Prepare blocks for forward pass
    offloader.prepare_block_devices_before_forward(blocks)
    
    # Simulate forward pass with offloading
    input_tensor = torch.randn(1, 10, device=device)
    x = input_tensor
    
    for i, block in enumerate(blocks):
        # Wait for the current block to be ready
        offloader.wait_for_block(i)
        
        # Process through the block
        x = block(x)
        
        # Schedule moving weights for future blocks
        offloader.submit_move_blocks(blocks, i)
    
    # Verify we get a valid output
    assert x.shape == (1, 10)
    assert not torch.isnan(x).any()


# Error handling tests