from accelerate import PartialState
from accelerate.utils import gather_object
import os
import torch

# Print CUDA information
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"Number of CUDA devices: {torch.cuda.device_count()}")
    print(f"Current CUDA device: {torch.cuda.current_device()}")
    print(f"CUDA device name: {torch.cuda.get_device_name()}")

# Initialize distributed state first
distributed_state = PartialState()

# Create some test data
prompts = [str(i) for i in range(4)]  # ["0", "1", "2", "3"]
batch_size = 2

# Split into batches
tokenized_prompts = [prompts[i : i + batch_size] for i in range(0, len(prompts), batch_size)]

completions_per_process = []
with distributed_state.split_between_processes(tokenized_prompts, apply_padding=True) as batched_prompts:
    for batch in batched_prompts:
        generated_text = [f"{distributed_state.device}: {t}" for t in batch]
        completions_per_process.extend(generated_text)

# Gather results from all processes
completions_gather = gather_object(completions_per_process)
print(completions_gather)

CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc-per-node 1 test.py

CUDA_VISIBLE_DEVICES="" torchrun --nproc-per-node 2 test.py