How to run inference of a (very) large model across mulitple GPUs ? about candle

[EricLBuehler]

Please see the llama multiprocess example. The multi-GPU inference is used to create parellelized linear layers:

candle/candle-examples/examples/llama_multiprocess/model.rs

Lines 293 to 308 in f48c07e

	fn load(vb: VarBuilder, cache: &Cache, cfg: &Config, comm: Rc<Comm>) -> Result<Self> {
	let qkv_proj = TensorParallelColumnLinear::load_multi(
	vb.clone(),
	&["q_proj", "k_proj", "v_proj"],
	comm.clone(),
	)?;
	let o_proj = TensorParallelRowLinear::load(vb.pp("o_proj"), comm.clone())?;
	Ok(Self {
	qkv_proj,
	o_proj,
	num_attention_heads: cfg.num_attention_heads / comm.world_size(),
	num_key_value_heads: cfg.num_key_value_heads / comm.world_size(),
	head_dim: cfg.hidden_size / cfg.num_attention_heads,
	cache: cache.clone(),
	})
	}

[b0xtch]

Please see the llama multiprocess example. The multi-GPU inference is used to create parellelized linear layers:

candle/candle-examples/examples/llama_multiprocess/model.rs

Lines 293 to 308 in f48c07e

	fn load(vb: VarBuilder, cache: &Cache, cfg: &Config, comm: Rc<Comm>) -> Result<Self> {
	let qkv_proj = TensorParallelColumnLinear::load_multi(
	vb.clone(),
	&["q_proj", "k_proj", "v_proj"],
	comm.clone(),
	)?;
	let o_proj = TensorParallelRowLinear::load(vb.pp("o_proj"), comm.clone())?;
	Ok(Self {
	qkv_proj,
	o_proj,
	num_attention_heads: cfg.num_attention_heads / comm.world_size(),
	num_key_value_heads: cfg.num_key_value_heads / comm.world_size(),
	head_dim: cfg.hidden_size / cfg.num_attention_heads,
	cache: cache.clone(),
	})
	}

That example is for a single node. How about multiple nodes? Can we just run the example with mpirun -n 2 --hostfile ../../hostfile target/release/llama_multiprocess 2 2000

Update:

I guess I must modify the code to support the world rank for MPI. I think sticking to NCCL as a backend might be better, but then is there support in Cudarc for cross-node communication?

Found this library https://github.com/oddity-ai/async-cuda