[train] Check when `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` can be set

[CharlieFRuan]

It is useful for saving memory consumption. But we need to make sure it does not break things.

From my understanding it can be safely set for non-colocated mode on the training node. For colocated setting, we can set it dynamically since I think it doesn't work well with CUDA IPC weight sync

[SumanthRH]

From my understanding it can be safely set for non-colocated mode on the training node. For colocated setting, we can set it dynamically since I think it doesn't work well with CUDA IPC weight sync

Actually expandable segments is not compatible with vLLM's cumem allocator. So one option is to set it just for the training worker processes.