parallel.py

import os
import torch
import gc
import logging


def auto_parallel(args):
    model_size = args.model_path.split("-")[-1]
    if model_size.endswith("m"):
        model_gb = 1
    else:
        model_gb = float(model_size[:-1])
    if model_gb < 20:
        n_gpu = 1
    elif model_gb < 50:
        n_gpu = 4
    else:
        n_gpu = 8
    args.parallel = n_gpu > 1
    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
    if isinstance(cuda_visible_devices, str):
        cuda_visible_devices = cuda_visible_devices.split(",")
    else:
        cuda_visible_devices = list(range(8))
    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
        [str(dev) for dev in cuda_visible_devices[:n_gpu]]
    )
    logging.debug("CUDA_VISIBLE_DEVICES: ", os.environ["CUDA_VISIBLE_DEVICES"])
    return cuda_visible_devices