Add parallelism_auto flag to automatically set dp, tp and micro batch size

[pefontana]

Adds --parallelism-auto flag to automatically detect optimal dp, tp, and micro_batch_size based on the model and GPU hardware.

1. Detects GPU type via nvidia-smi (with fallback to /proc/driver/nvidia/gpus/*/information for containers)
2. Fetches parallelism_data.json from the model's HuggingFace repo
3. Looks up config: GPU type → num GPUs → {dp, tp, micro_batch_size}

  {
    "H100": {
      "1": { "dp": 1, "tp": 1, "micro_batch_size": 4 },
      "8": { "dp": 4, "tp": 2, "micro_batch_size": 4 }
    },
    "H200": {
      "8": { "dp": 8, "tp": 1, "micro_batch_size": 8 }
    }
  }

We considered sharing the parallelism config via P2P but it added unnecessary complexity—dp/tp/micro_batch_size are needed before model download begins, requiring a separate request/wait flow. Since the file is ~200 bytes, fetching from HF is fast and keeps the code simple.

Note:

I think we can merge this PR now, we should only take into account that if the run is private, is posible that some trainers dont have access to HuggingFace, so I dont recommend using this flag for private runs now
Once #455 is merged, all trainers must have access to HF/GCP, so it wont be a problem

[entropidelic]

small nit: maybe we can rename this variable to something better?

[pefontana]

done!

[arilotter]

have you considered using the nvml_wrapper crate instead of shelling out / reading /proc/fs stuff? we can grab gpu count from there too 🤷 and assert that they're all the same GPU for sanity checking :D

we use this in some of the metrics stuff already -
it would be something like:

use nvml_wrapper::Nvml;

#[derive(Debug)]
struct GpuInfo {
    name: String,
    device_count: u32,
}

fn get_gpu_info() -> anyhow::Result<GpuInfo> {
    let nvml = Nvml::init()?;
    let device_count = nvml.device_count()?;

    if device_count == 0 {
        anyhow::bail!("No GPUs found!");
    }

    let mut gpu_names = Vec::new();

    for i in 0..device_count {
        let device = nvml.device_by_index(i)?;
        gpu_names.push(device.name()?);
    }

    let first_name = &gpu_names[0];
    if !gpu_names.iter().all(|name| name == first_name) {
        anyhow::bail!(
            "All GPUs must be of the same type, but we have mismatching names: {:?}",
            gpu_names
        );
    }

    Ok(GpuInfo {
        name: gpu_names.pop().unwrap(),
        device_count,
    })
}

[pefontana]

Thanks Ari!
I add it and works

[pefontana]

I just changed the let device_count = nvml.device_count()?; because was not taking into account CUDA_VISIBLE_DEVICES env var