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

architectures/centralized/client/src/app.rs

@@ -131,6 +131,7 @@ pub async fn build_app(
     .await?;
 
     let state_options: RunInitConfig<ClientId, ClientId> = RunInitConfig {
+        parallelism_auto: p.parallelism_auto,
         data_parallelism: p.data_parallelism,
         tensor_parallelism: p.tensor_parallelism,
         micro_batch_size: p.micro_batch_size,

architectures/decentralized/solana-client/src/app.rs

@@ -114,6 +114,7 @@ pub async fn build_app(
 
     let state_options: RunInitConfig<psyche_solana_coordinator::ClientId, NetworkIdentity> =
         RunInitConfig {
+            parallelism_auto: p.parallelism_auto,
             data_parallelism: p.data_parallelism,
             tensor_parallelism: p.tensor_parallelism,
             micro_batch_size: p.micro_batch_size,

nix/lib.nix

@@ -20,7 +20,8 @@ let
     || (builtins.match ".*tests/fixtures/.*$" path != null)
     || (builtins.match ".*.config/.*$" path != null)
     || (builtins.match ".*local-dev-keypair.json$" path != null)
-    || (builtins.match ".*shared/client/src/state/prompt_texts/index\\.json$" path != null);
+    || (builtins.match ".*shared/client/src/state/prompt_texts/index\\.json$" path != null)
+    || (builtins.match ".*shared/client/src/parallelism_data\\.json$" path != null);
 
   src = lib.cleanSourceWith {
     src = ../.;

psyche-book/src/enduser/create-run.md

@@ -86,11 +86,46 @@ run-manager create-run \
 
 At this point, your run has been successfully created.
 
+### Adding parallelism configuration (recommended)
+
+If your run uses a model that is not already in the [parallelism lookup table](https://github.com/PsycheFoundation/psyche/blob/main/shared/client/src/parallelism_data.json), it's recommended to add the optimal parallelism configuration for your model and target GPU hardware. This allows clients to use `PARALLELISM_AUTO=true` for automatic configuration.
+
+#### Option 1: Add to your model's HuggingFace repo (preferred)
+
+Add a `parallelism_data.json` file directly to your model's HuggingFace repository. The client will automatically fetch this configuration at runtime - no Psyche rebuild required.
+
+```json
+{
+	"your-org/your-model": {
+		"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 }
+		}
+	}
+}
+```
+
+This is the preferred approach because:
+
+- No PR to Psyche required
+- No Docker image rebuild needed
+- Model creators manage their own configuration
+- Changes take effect immediately
+
+#### Option 2: Add to Psyche's compiled table
+
+Alternatively, open a PR to add your model to the [compiled parallelism table](https://github.com/PsycheFoundation/psyche/blob/main/shared/client/src/parallelism_data.json). This is useful for widely-used models that should have default configurations. The format is the same as above.
+
+**Lookup order**: The client first tries to fetch config from the model's HuggingFace repo, then falls back to the compiled table with a warning if not found.
+
 ### Initializing configuration
 
 Initially, the run will not have any configuration defined and will remain paused, so no clients can join yet.
 
-To set the run configuration, you’ll need to provide mostly the same parameters as when creating the run, along with the path to a `config.toml` file that follows the [run config schema](./run-config.md).
+To set the run configuration, you'll need to provide mostly the same parameters as when creating the run, along with the path to a `config.toml` file that follows the [run config schema](./run-config.md).
 
 ```bash
 run-manager update-config \

psyche-book/src/enduser/join-run.md

@@ -93,19 +93,29 @@ though you might need to.
 
 **`NVIDIA_DRIVER_CAPABILITIES`** - An environment variable that the NVIDIA Container Toolkit uses to determine which compute capabilities should be provided to your container. It is recommended to set it to 'all', e.g. `NVIDIA_DRIVER_CAPABILITIES=all`.
 
+**`PARALLELISM_AUTO`** - Set to `true` to automatically detect optimal parallelism settings based on the model and your GPU hardware.
+
+- When enabled, the client will look up the best `DATA_PARALLELISM`, `TENSOR_PARALLELISM`, and `MICRO_BATCH_SIZE` values from a [built-in configuration table](https://github.com/PsycheFoundation/psyche/blob/main/shared/client/src/parallelism_data.json)
+- Your model and GPU hardware combination must be present in the table
+- This is the recommended option for most users
+- If set, manual parallelism settings below will be ignored
+
 **`DATA_PARALLELISM`** - Number of GPUs to distribute training data across.
 
 - If you have multiple GPUs, you can set this to 2, 4, etc. to speed up training
 - If you have 1 GPU, set this to `1`
+- Ignored if `PARALLELISM_AUTO=true`
 
 **`TENSOR_PARALLELISM`** - Number of GPUs to distribute the model across, this lets you train a model you can't fit on one single GPU.
 
 - If you have 1 GPU, set this to `1`
 - If your have `n` GPUs you can distribute the model across all of them by setting it to `n`.
+- Ignored if `PARALLELISM_AUTO=true`
 
 **`MICRO_BATCH_SIZE`** - Number of samples processed per GPU per training step
 
 - Set as high as your GPU memory allows
+- Ignored if `PARALLELISM_AUTO=true`
 
 **`AUTHORIZER`** - The Solana address that authorized your wallet to join this run
 

shared/client/src/cli.rs

@@ -112,6 +112,10 @@ pub struct TrainArgs {
     #[clap(long, env, value_parser = parse_trim_quotes)]
     pub run_id: String,
 
+    /// Auto-detect parallelism settings from lookup table based on model and GPU count
+    #[clap(long, env)]
+    pub parallelism_auto: bool,
+
     #[clap(long, default_value_t = 1, env)]
     pub data_parallelism: usize,
 

shared/client/src/lib.rs

@@ -1,6 +1,7 @@
 mod cli;
 mod client;
 mod fetch_data;
+pub mod parallelism_lookup;
 mod protocol;
 mod state;
 mod tui;

shared/client/src/parallelism_data.json

@@ -0,0 +1,40 @@
+{
+	"_doc (HuggingFace repo)": {
+		"gpu type from nvidia-smi": {
+			"numember of gpus": {
+				"dp": 0,
+				"tp": 0,
+				"micro_batch_size": 0
+			}
+		}
+	},
+	"emozilla/llama2-20m-init": {
+		"H100": {
+			"1": { "dp": 1, "tp": 1, "micro_batch_size": 1 },
+			"8": { "dp": 1, "tp": 1, "micro_batch_size": 1 }
+		}
+	},
+	"TinyLlama/TinyLlama-1.1B-Chat-v0.4": {
+		"H100": {
+			"1": { "dp": 1, "tp": 1, "micro_batch_size": 8 },
+			"8": { "dp": 8, "tp": 1, "micro_batch_size": 8 }
+		}
+	},
+	"NousResearch/Meta-Llama-3.1-8B": {
+		"H100": {
+			"1": { "dp": 1, "tp": 1, "micro_batch_size": 1 },
+			"8": { "dp": 8, "tp": 1, "micro_batch_size": 4 }
+		}
+	},
+	"deepseek-ai/DeepSeek-V2-Lite": {
+		"H100": {
+			"1": { "dp": 1, "tp": 1, "micro_batch_size": 1 },
+			"8": { "dp": 4, "tp": 2, "micro_batch_size": 2 }
+		}
+	},
+	"NousResearch/Hermes-4-70B": {
+		"H100": {
+			"8": { "dp": 1, "tp": 8, "micro_batch_size": 1 }
+		}
+	}
+}

shared/client/src/parallelism_lookup.rs

@@ -0,0 +1,125 @@
+use anyhow::Result;
+use hf_hub::{Repo, RepoType};
+use serde::Deserialize;
+use std::collections::HashMap;
+use std::process::Command;
+use tracing::{info, warn};
+
+const PARALLELISM_DATA: &str = include_str!("parallelism_data.json");
+const REMOTE_CONFIG_FILENAME: &str = "parallelism_data.json";
+
+#[derive(Debug, Clone, Copy, Deserialize)]
+pub struct ParallelismConfig {
+    pub dp: usize,
+    pub tp: usize,
+    pub micro_batch_size: usize,
+}
+
+// Table format: model -> gpu_type -> num_gpus -> config
+type Table = HashMap<String, HashMap<String, HashMap<String, ParallelismConfig>>>;
+
+fn get_gpu_type() -> String {
+    // Try nvidia-smi first
+    let raw_gpu_name = Command::new("nvidia-smi")
+        .args(["--query-gpu=name", "--format=csv,noheader"])
+        .output()
+        .ok()
+        .and_then(|o| String::from_utf8(o.stdout).ok())
+        .and_then(|s| s.lines().next().map(|l| l.trim().to_string()))
+        .filter(|s| !s.is_empty())
+        // Fallback: read from /proc/driver/nvidia (works in containers without nvidia-smi)
+        .or_else(|| {
+            std::fs::read_dir("/proc/driver/nvidia/gpus")
+                .ok()?
+                .filter_map(|e| e.ok())
+                .next()
+                .and_then(|entry| {
+                    let info_path = entry.path().join("information");
+                    std::fs::read_to_string(info_path).ok()
+                })
+                .and_then(|content| {
+                    content
+                        .lines()
+                        .find(|line| line.starts_with("Model:"))
+                        .map(|line| line.trim_start_matches("Model:").trim().to_string())
+                })
+        })
+        .unwrap_or_default();
+
+    // Normalize GPU name to match table keys
+    if raw_gpu_name.to_uppercase().contains("H200") {
+        "H200".to_string()
+    } else if raw_gpu_name.to_uppercase().contains("H100") {
+        "H100".to_string()
+    } else {
+        raw_gpu_name
+    }
+}
+
+/// Try to load parallelism config from the model's HuggingFace repo
+fn load_from_model_repo(model_repo_id: &str) -> Option<Table> {
+    let token = std::env::var("HF_TOKEN").ok();
+
+    let api = hf_hub::api::sync::ApiBuilder::new()
+        .with_token(token)
+        .build()
+        .ok()?
+        .repo(Repo::new(model_repo_id.to_string(), RepoType::Model));
+
+    let path = api.get(REMOTE_CONFIG_FILENAME).ok()?;
+    let content = std::fs::read_to_string(path).ok()?;
+    serde_json::from_str(&content).ok()
+}
+
+/// Lookup config in a table
+fn lookup_in_table(
+    table: &Table,
+    model_repo_id: &str,
+    gpu_type: &str,
+    num_gpus: usize,
+) -> Option<ParallelismConfig> {
+    table
+        .get(model_repo_id)
+        .and_then(|g| g.get(gpu_type))
+        .and_then(|n| n.get(&num_gpus.to_string()))
+        .copied()
+}
+
+/// Load the compiled parallelism table
+fn load_compiled_table() -> Result<Table> {
+    serde_json::from_str(PARALLELISM_DATA)
+        .map_err(|e| anyhow::anyhow!("Failed to parse compiled parallelism data: {}", e))
+}
+
+pub fn lookup(model_repo_id: &str) -> Result<ParallelismConfig> {
+    let num_gpus = tch::Cuda::device_count() as usize;
+    let gpu_type = get_gpu_type();
+    info!("Detected {} x {} GPU(s)", num_gpus, gpu_type);
+
+    // Try model's own config first
+    if let Some(table) = load_from_model_repo(model_repo_id) {
+        if let Some(config) = lookup_in_table(&table, model_repo_id, &gpu_type, num_gpus) {
+            info!(
+                "Using parallelism config from model repo '{}'",
+                model_repo_id
+            );
+            return Ok(config);
+        }
+    }
+
+    // Fall back to compiled table
+    warn!(
+        "No parallelism config found in model repo '{}', using compiled defaults",
+        model_repo_id
+    );
+
+    let table = load_compiled_table()?;
+    lookup_in_table(&table, model_repo_id, &gpu_type, num_gpus).ok_or_else(|| {
+        anyhow::anyhow!(
+            "No config for {} x {} with model '{}'",
+            num_gpus,
+            gpu_type,
+            model_repo_id
+        )
+    })
+}