Remove tms while preserving offloading

[Risc-lt]

This PR removes torch memory saver while preserving offloading function. The overview of the offloading mechanism is

# Initialization
x = torch.empty(1000, 1000, device="cuda")

# Release memory
x.storage().resize_(0)
torch.cuda.empty_cache()

# Realloc when onloading
# Since we don't need the last value, we can skip the h2d and just do the allocation
x.storage().resize_(x.numel())

After testing on 1to1 config, the effect is

[RDMA] Before offloading model replica: {'gpu': '0', 'total_GB': 139.8, 'free_GB': 57.41, 'used_GB': 82.39}

[RDMA] After offloading model replica: {'gpu': '0', 'total_GB': 139.8, 'free_GB': 66.13, 'used_GB': 73.67}

[Risc-lt]

Current profiling shows that registering results in 900~1k ms overhead and unregister is ~300 ms. We need to find some way to pipeline the process. cc @JD-ETH @JensenFire

[JD-ETH]

let's discuss a bit first

[JD-ETH]

either make it private and access via a property, or just make it public member

[Risc-lt]

sry for this typo, I'll correct it

[JD-ETH]

don't we have _active_tasks already?

[JD-ETH]

we should rely on _queue.join() for the eventual finish check

[JD-ETH]

is there a way to guarantee the mapping is exact? calling memory_snapshot can be expensive, no?

I think its best if we store the engine_param -> [memory, offset] mapping at registration time.

[Risc-lt]

Exactly, the best way here is to return the registered memory region address in function register_memory_region_v2 (sglang side). Then we can directly pass these addresses to TE.unregister. Current implementation is to mock register_memory_region_v2 on traning side.

[JD-ETH]

for now let's pre-commit and fix the private member issue --- we can merge first, and work on correctness as an initial step

[Risc-lt]

Thanks for review! I'll do the pre-commit and resolve the problems commented.