[Enhancement]: use set_cuda_visible_devices instead of original device_map

[xin3he]

Problem Description

This usage 0.9 is not stable, By setting CUDA_VISIBLE_DEVICES, `get_balanced_memory` will distribute the model evenly across all visible devices.

Reproduction Steps

Currently, we have set_cuda_visible_devices, but only used in eval mode

Environment Information

No response

Error Logs

Additional Context

No response

[xin3he]

cc @wenhuach21

[xin3he]

To avoid breaking change, Added warning to set CUDA_VISIBLE_DEVICES when using get_balanced_memory