feat: Multi-device training

[bradhilton]

No description provided.

[Danau5tin]

Hi @bradhilton , thanks for notifying me!
Yes I would love to try it out, I really want to train Qwen3-32B, and normally would use deepspeed zero 3 for that and split the inference and training over multiple GPUs.

Am I right in saying (quick look), that ART doesn't split GPUs between training and inference. Instead, it uses all GPUs sequentially?
So if you have 8 GPUs, all 8 are used for inference, then all 8 for training, cycling back and forth as needed?

Do you think an (8x in practice) A100 (80GB) is fine for Qwen3-32B?

[bradhilton]

@Danau5tin yep, GPUs are colocated. This means you get high GPU utilization without having to figure out the optimal balance between inference and training GPUs.

If there is enough interest from those who are more concerned about absolute best wall time, whatever the cost, than utilization or those who are willing to find the optimal number of training and inference GPUs themselves to maximize utilization, then maybe we will add support for separate deployments in the future.

I've been using 8xH200s for Qwen3-32B which have about double the memory, but if your sequence length is not too long and you lower vLLM's GPU utilization you may be okay. You can also try enabling activation offloading to save additional memory.

Here's how you initialize your trainable model to use the multi-device backend at the moment:

model = art.TrainableModel(
    name="model-name",
    project="project-name",
    base_model="Qwen/Qwen3-32B",
    _internal_config=art.dev.InternalModelConfig(
        engine_args=art.dev.EngineArgs(
            tensor_parallel_size=torch.cuda.device_count(), # need to specify the tensor parallel size for vLLM, at least for now
            gpu_memory_utilization=0.9 # defaults to 0.9, but you will probably need to lower this to avoid OOM errors
        ),
        torchtune_args=art.dev.TorchtuneArgs(
            model="qwen3_32b", # needed for torchtune model initialization
            model_type="QWEN3", # needed for torchtune checkpointing
            async_weight_syncing=False # defaults to False, changes the learning guaruntees to enable but speeds up training
            enable_activation_offloading=False # defaults to False, enable to save VRAM for a small speed hit
        ),
    ),
)

[Danau5tin]

Hey @bradhilton, so I migrated my project over to ART, tried out the code, and all looked great!
Then I realised the model wasn't training.

Here are some logs I have:

 Step 3: Generating trajectories: 100%|██████████████████████████████| 4/4 [01:51<00:00, 27.99s/it, reward=0.104, turns=3, test_reward=0, judge_reward=0.297, prompt_tokens=2.85e+3, total_completion_tokens=9386.0]
  
2025-07-03 18:29:15,391 - src.train_with_art - INFO - Collected 4 which have rewards: [0.147, 0.0875, 0.1295, 0.0525]. Stnd dev: 0.04257, prompt_tokens=2.85e+3, total_completion_tokens=9386.0]
  
Skipping tuning as there is no suitable data. This can happen when all the trajectories in the same group have the same reward and thus no advantage to train on.

Created with this code:

        all_rewards = [
            trajectory.reward for group in train_groups for trajectory in group.trajectories
        ]
        logger.info(
            f"Collected {len(all_rewards)} which have rewards: {all_rewards}. Stnd dev: {torch.std(torch.tensor(all_rewards)):.4f}"
        )

Not sure where that warning at the bottom comes from, but it comes up no matter what the rewards in the group are, and I can see from other logs that the trajectories are actually happening as expected.

So to see if it was because of the multi-gpu setup, I switched to single GPU (removed this internal config below - just commented it out):

    internal_config = art.dev.InternalModelConfig(
        engine_args=art.dev.EngineArgs(
            tensor_parallel_size=torch.cuda.device_count(),
            gpu_memory_utilization=0.9,
        ),
        torchtune_args=art.dev.TorchtuneArgs(
            model=torchtune_model,
            model_type=model_type,
        ),
    )

When I switched to the standard way of training on a single GPU, it just seemed to hang on the first rollout. So not sure what is happening there either because the rollouts worked fine with the multi-gpu setup.

It has been a long day, so I have only tried the single GPU setup - with Unsloth - once before I wrote this! Perhaps If i tried again it would work.

I am using Qwen 0.6B for reference (Just to try and get a successful run, then I will increase model params), on 2x A100 GPUs.

[bradhilton]

@Danau5tin make sure that the trajectories in a group have differing rewards. Standard deviation should be calculated in group and not between groups because that does not affect the advantage calculation. If you want more timely help, feel free to join the discord!