[BUG]Resolving OOM Issues in ConcurrenDistributed Inference of 111B Teacher Model and Distributed Training of 8B Student Model on Multi-Node H200 GPUs

Hello, first of all, apologies if this information is already written somewhere. I've searched through documentation and issues but haven't found a clear answer to my question.

I have access to 2 to 4 nodes (16 to 32 GPUs) with 8x140GB H200 GPUs each. My goal is to perform large-scale distributed inference using a large 111B-parameter model (CohereLabs/c4ai-command-a-03-2025) to perform online knowledge distillation (soft-logit based) from the 111B Teacher model to a smaller 8B Student model (CohereLabs/c4ai-command-r7b-12-2024). Is there a way to simultaneously run Teacher distributed inference and Student distributed training on multiple nodes?

Currently, it seems that distributed inference and distributed training across multiple nodes are not functioning correctly, as the memory usage appears similar when using a single node with 8x140GB H200 GPUs and when using multiple nodes with 24x140GB H200 GPUs.

I'm still struggling with multi-node distribution even after trying both deepspeed.init_inference and deepspeed.initialize. Specifically, for Teacher models larger than 111B, a single node results in OOM errors, which is why I'm attempting to use multiple nodes. However, even with 4 or more nodes (32 H200 GPUs) using the code below, OOM errors persist.

I've been trying to use DeepSpeed-Inference, but from my understanding, before using deepspeed.init_inference, I have to load the entire model onto the GPU or CPU. This process always leads to an Out-of-Memory (OOM) error because the Teacher model is loaded onto each single node, and inter-node distribution doesn't seem to be happening.

I would be very grateful if someone could explain what would be most suitable for my use case. A minimal working example would be the icing on the cake. Surely, if the Open LLM Leaderboard shows that online knowledge distillation (soft-logit) is possible with large models exceeding 111B, there must be a straightforward way to achieve what I want, but I'm unsure how everyone else does it.

For reference, below is the script I'm currently working with:

deepspeed --num_nodes 2 --num_gpus 8 \ --hostfile $HOSTFILE \ --master_addr $MASTER_ADDR \ --master_port=62535 \ train.py \ --teacher CohereLabs/c4ai-command-a-03-2025 \ --student CohereLabs/c4ai-command-r7b-12-2024 \ --epochs 1 --batch_size 1 --seq_len 4096 --temperature 1.0 --max_samples 150 --lr 1e-6 2>&1 | tee -a "./train.log"

import torch.distributed as dist
import os, math, argparse, warnings, torch, random, multiprocessing as mp
from datasets import load_dataset, concatenate_datasets
from transformers import (AutoTokenizer, AutoModelForCausalLM,
                          PreTrainedTokenizerBase)
from torch.nn.utils.rnn import pad_sequence
import torch.nn.functional as F
from datetime import timedelta
from deepspeed.runtime.utils import see_memory_usage


os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ.setdefault("NCCL_ASYNC_ERROR_HANDLING", "1")
warnings.filterwarnings("ignore", category=UserWarning)
mp.set_start_method("spawn", force=True)

def get_args():
    p = argparse.ArgumentParser()
    p.add_argument("--teacher", default="")
    p.add_argument("--student", default="")
    p.add_argument("--dataset", default="")
    p.add_argument("--split", default="train")
    p.add_argument("--epochs", type=int, default=1)
    p.add_argument("--batch_size", type=int, default=1,
                   help="per-GPU micro-batch")
    p.add_argument("--seq_len", type=int, default=4096)
    p.add_argument("--temperature", type=float, default=1.0)
    p.add_argument("--lr", type=float, default=1e-6)
    p.add_argument("--max_samples", type=int, default=0,
                   help="0=1000 ")
    p.add_argument("--local_rank", type=int, default=-1,
               help="deepspeed/torch launcher GPU index")
    p.add_argument("--cache_path", default="")
    p.add_argument("--hf_token", default="")
    p = deepspeed.add_config_arguments(p)
    return p.parse_args()


def main():
    timeout_seconds = 3600 
    timeout_duration = timedelta(seconds=timeout_seconds)
    dist.init_process_group(
        backend="nccl",
        timeout=timeout_duration 
    )
    args = get_args()
    deepspeed.init_distributed()
    rank, world = deepspeed.comm.get_rank(), deepspeed.comm.get_world_size()
    device = torch.device("cuda", deepspeed.comm.get_local_rank())
    # Tokenizer 
    tokenizer = AutoTokenizer.from_pretrained(args.student,
                                        use_fast=True, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        
    # tokenizer token_id 
    tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids(tokenizer.eos_token)
    tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
    
    
    # Teacher (inference only)    
    teacher_model = AutoModelForCausalLM.from_pretrained(
        args.teacher, torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
        trust_remote_code=True, device_map=None, 
        cache_dir=args.cache_path,token=args.hf_token) 
    
    see_memory_usage("After load model", force=True)
    
    teacher_model.config.eos_token_id = tokenizer.eos_token_id
    teacher_model.config.pad_token_id = tokenizer.pad_token_id
        
    teacher_engine = deepspeed.init_inference(
        teacher_model,
        mp_size=world,
        dtype=torch.bfloat16,
        replace_with_kernel_inject=True, 
        replace_method="auto")
    
    see_memory_usage("After DS-inference init", force=True)

    teacher_engine.module.eval()
    teacher_engine.optimizer = None

    # Student
    student_model = AutoModelForCausalLM.from_pretrained(
        args.student, torch_dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
        trust_remote_code=True, cache_dir=args.cache_path,token=args.hf_token)
    
    student_model.config.eos_token_id = tokenizer.eos_token_id
    student_model.config.pad_token_id = tokenizer.pad_token_id

    # Dataset 
    ds = [
        load_dataset("Raphael21/LogicKor_Aug_small_v2", split="train", data_dir="v0.1.1", streaming=False)
    ]

    ds = concatenate_datasets(ds).select_columns(["messages"])
    total_samples = args.max_samples or len(ds)
    total_steps = args.epochs * math.ceil(total_samples / args.batch_size)
    
    # Deepspeed Config 
    ds_cfg = {
        "train_batch_size": args.batch_size * world,
        "gradient_accumulation_steps": 1,
        "bf16": {"enabled": True},
        "zero_optimization": {
            "stage": 3,  
            "stage3_max_live_parameters": 1e9,
            "stage3_prefetch_bucket_size": 5e8,
            "stage3_param_persistence_threshold": 1e4,
            "overlap_comm": True,
            "contiguous_gradients": True,
            "allgather_bucket_size": 5e8,
            "reduce_bucket_size": 5e8,
            "offload_optimizer": {
                "device": "cpu", 
                "pin_memory": True
            },
            "offload_param": {
                "device": "cpu", 
                "pin_memory": True
            },
        },
        "activation_checkpointing": {
            "partition_activations": True,
            "contiguous_memory_optimization": True
        },
        "optimizer": {
            "type": "AdamW",
            "params": {
                "lr": args.lr,
                "betas": [0.9, 0.999],
                "eps": 1e-8,
                "weight_decay": 0.01
            }
        },
       
        "scheduler": {
            "type": "WarmupLR",
            "params": {
                "warmup_min_lr": 0,
                "warmup_max_lr": args.lr,
                "warmup_num_steps": int(0.1 * total_steps)
            }
        }
    }


    student_engine, _, _, scheduler = deepspeed.initialize(
        args=args,
        model=student_model,
        model_parameters=student_model.parameters(),
        config=ds_cfg
    )
    
    
    if not hasattr(student_engine, "optimizer"):
        student_engine.optimizer = None
        
    # Debug Messages
    if rank == 0:
        print("Configured with ZeRO-3, total_steps:", total_steps)
        
    # Data Loader
    def preprocess_batch(examples):
        prompt_key = "prompt"
        messages_key = "messages"
        ignore_index = -100
        max_length = min(tokenizer.model_max_length, 4096)

        def get_tokens_from_chat_template(messages_dict_or_str, add_gen_prompt, max_len=max_length):
            
            tokens = tokenizer.apply_chat_template(
                messages_dict_or_str,
                tokenize=True,             
                add_generation_prompt=add_gen_prompt,
                truncation=True,
                max_length=max_len,
                padding="do_not_pad",      
                return_tensors=None,
            )
           
            return tokens, [1] * len(tokens) 
        
        results = {
            "input_ids": [],
            "attention_mask": [],
            "labels": [],
            "prompts": [],
            "prompt_attention_mask": [],
        }

        prompts = examples.get(prompt_key, [None] * len(examples[messages_key]))
        for message, prompt in zip(examples[messages_key], prompts):
            if prompt is None:
                prompt_messages = message[:-1] 
                prompt_ids, prompt_attn = get_tokens_from_chat_template(prompt_messages, add_gen_prompt=True)
            else:
                prompt_ids, prompt_attn = get_tokens_from_chat_template(prompt_text, add_gen_prompt=True) 
                
            input_ids, attn_mask = get_tokens_from_chat_template(message, add_gen_prompt=False)
            label = [ignore_index] * len(input_ids) 
            start_idx = len(prompt_ids)

            if start_idx < len(input_ids):
                for i in range(start_idx, len(input_ids)):
                    label[i] = input_ids[i]            

            results["input_ids"].append(input_ids)
            results["attention_mask"].append(attn_mask)
            results["labels"].append(label)
            results["prompts"].append(prompt_ids)
            results["prompt_attention_mask"].append(prompt_attn)
            
        return results

    def chatml_collate_fn(batch, pad_token_id=0, ignore_index=-100, max_length=4096):
        def pad_and_truncate(seqs, pad_val, max_length):
            return torch.stack([
                torch.tensor(seq[:max_length] + [pad_val] * (max_length - len(seq)), dtype=torch.long)
                for seq in seqs
            ])

        fields = ["input_ids", "attention_mask", "labels", "prompts", "prompt_attention_mask"]
        pad_values = [pad_token_id, 0, ignore_index, pad_token_id, 0]

        return {
            field: pad_and_truncate([ex[field] for ex in batch], pad_val, max_length)
            for field, pad_val in zip(fields, pad_values)
        }


    if args.max_samples:
        ds = ds.select(range(args.max_samples))

    ds = ds.map(preprocess_batch, batched=True)

    loader = torch.utils.data.DataLoader(
        ds,
        batch_size=args.batch_size,
        shuffle=True,
        pin_memory=True,
        collate_fn=lambda x: chatml_collate_fn(
            x,
            pad_token_id=tokenizer.pad_token_id,
            ignore_index=-100,
            max_length=args.seq_len  
        )
    )
   
    T = args.temperature

    for epoch in range(args.epochs):
        for step, batch in enumerate(loader):
            prompt_lengths_batch = batch["prompt_attention_mask"].sum(dim=1).cpu().tolist()
            prompt_lengths_tensor = torch.tensor(prompt_lengths_batch, device=device, dtype=torch.long)
            
            input_ids = batch["input_ids"].to(device)
            attn = batch["attention_mask"].to(device)
            labels_batch = batch["labels"].to(device)

            with torch.no_grad():
                teacher_logits = teacher_engine.module( 
                    input_ids=input_ids,
                    attention_mask=attn,
                    use_cache=False
                ).logits

            student_logits = student_engine(
                input_ids=input_ids,
                attention_mask=attn,
                use_cache=False
            ).logits
            
            
            if rank == 0:
                sample_idx_to_inspect = 0
                original_input_ids = batch["input_ids"][sample_idx_to_inspect].cpu().tolist()
                original_labels_list = batch["labels"][sample_idx_to_inspect].cpu().tolist() 

                # Student Model (argmax)
                student_predictions_ids = student_logits[sample_idx_to_inspect].argmax(dim=-1).cpu().tolist()
                decoded_student_predictions = [tokenizer.decode([t], skip_special_tokens=False) for t in student_predictions_ids]

                # Teacher Model (argmax)
                teacher_predictions_ids = teacher_logits[sample_idx_to_inspect].argmax(dim=-1).cpu().tolist()
                decoded_teacher_predictions = [tokenizer.decode([t], skip_special_tokens=False) for t in teacher_predictions_ids]

                print(f"Decoded Student Predictions: {''.join(decoded_student_predictions[:100])} ...")
                print(f"Decoded Teacher Predictions: {''.join(decoded_teacher_predictions[:100])} ...")

            shifted_student_logits = student_logits[:, :-1, :] 
            shifted_teacher_logits = teacher_logits[:, :-1, :]
            shifted_labels = labels_batch[:, 1:] 
            shifted_attention_mask = attn[:, 1:] 
            
            current_seq_len = shifted_labels.size(1) 
            response_mask = torch.zeros_like(shifted_labels, dtype=torch.bool)

            for i in range(args.batch_size):
                start_response_idx_in_shifted = prompt_lengths_tensor[i] - 1
                start_response_idx_in_shifted = max(0, start_response_idx_in_shifted)
                if start_response_idx_in_shifted < current_seq_len:
                    response_mask[i, start_response_idx_in_shifted:] = True
                    
            shifted_attention_mask = shifted_attention_mask & response_mask        
                 
            # Apply temperature scaling
            student_logits_scaled = shifted_student_logits / args.temperature
            teacher_logits_scaled = shifted_teacher_logits / args.temperature

            # Compute log probabilities for student and probabilities for teacher
            student_log_probs = F.log_softmax(student_logits_scaled, dim=-1)
            teacher_log_probs = F.log_softmax(teacher_logits_scaled, dim=-1)
            
            kd_loss = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True)
            
            kd_loss_per_token = kd_loss.sum(dim=-1) 
            
            valid_labels_mask = (shifted_labels != -100) 
            combined_mask = shifted_attention_mask & valid_labels_mask 
            masked_kd_loss = kd_loss_per_token * combined_mask
            
            
            num_valid_tokens = combined_mask.sum()
            if num_valid_tokens > 0:
                kd_loss = masked_kd_loss.sum() / num_valid_tokens
            else: 
                kd_loss = torch.tensor(0.0, device=device, requires_grad=True) 
            
            # Cross-Entropy Loss 
            ce_loss = F.cross_entropy(
                shifted_student_logits.view(-1, shifted_student_logits.size(-1)), # (B*S, V)
                shifted_labels.view(-1),
                ignore_index=-100 
            )
            if ce_loss.numel() == 0 or torch.isnan(ce_loss):
                ce_loss = torch.tensor(0.0, device=device, requires_grad=True)
                
            alpha = 0.5 
            total_loss = alpha * ce_loss + (1 - alpha) * kd_loss

            student_engine.backward(total_loss)
            student_engine.step()
            
            # empty cache
            torch.cuda.empty_cache()

            if rank == 0 and step % 1 == 0:
                print(f"[Epoch {epoch}][{step}/{len(loader)}] total_loss = {total_loss.item():.4f}, ce_loss = {ce_loss.item():.4f}, kd_loss = {kd_loss.item():.4f}")
               
    # Save Checkpoint 
    student_engine.save_checkpoint("./save_checkpoint")

if __name__ == "__main__":
    main()```

Start with the train.py command and inspect the teacher model loading before deepspeed.init_inference, then compare the reported memory behavior across node counts. Review the DeepSpeed inference and training initialization paths shown in the script. Done means identifying whether the multi-node setup is supported as written and documenting a reproducible configuration or minimal example that avoids the reported OOM.