GPU Node Crash During Training on H100 GPUs

[PhoenixAI0]

Hello,

I'm experiencing an issue while training a LLM using the Prime framework on H100 GPUs. The training script consistently crashes the entire GPU node after printing "[INFO] [Rank 0] starting training." This issue is reproducible across multiple attempts and does not generate a standard CUDA Out-of-Memory (OOM) error.

Issues Observed:

SSH access to the node is lost after a while post "starting training" message.

02:03:36 [INFO] [Rank 0] [Rank 0] starting training
Connection to 1aj9oe98idds3evvlu75f45vq0.ingress.h100.hou.val.akash.pub closed by remote host.
Connection to 1aj9oe98idds3evvlu75f45vq0.ingress.h100.hou.val.akash.pub closed.

No anomalies were observed on the GPU rental dashboard.
When I attempt to reconnect via SSH, the fingerprint seems to have changed, outputting something like this:

@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@    WARNING: REMOTE HOST IDENTIFICATION HAS CHANGED!     @
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
IT IS POSSIBLE THAT SOMEONE IS DOING SOMETHING NASTY!
Someone could be eavesdropping on you right now (man-in-the-middle attack)!
Host key for [1aj9oe98idds3evvlu75f45vq0.ingress.h100.hou.val.akash.pub]:31750 has changed and you have requested strict checking.
Host key verification failed.

After I use ssh-keygen and reconnect, the system appears reset (training folder and packages are gone).

Environment:

GPU: H100 (SXM5 4xH100s rented through Akash on the Prime Intellect platform)
Model: 150M parameter LLaMA-like architecture
Dataset: Fineweb-Edu
Operating System: Ubuntu, using the Pytorch 2.2 Cluster Base Image
Python Version: Python 3.10.12
Config Parameters:

name_model = "150M"
type_model = "llama2"
project = "Phoenix_Test"
run_id = "phoenix-test"

[train]
micro_bs = 16
torch_compile = true
ac_ckpt = 1
reshard_after_forward = true
sequence_packing = true
attn_fn = "flash"

[optim]
batch_size = 1048576
lr = 6e-4
weight_decay = 0.1
adam_betas1 = 0.9
adam_betas2 = 0.95
sched_type = "cosine"
warmup_steps = 700
total_steps = 18865

[data]
dataset_name_or_paths = "edu_fineweb100B/"
val_dataset_name_or_paths = "edu_fineweb_val_000000.parquet"
seq_length = 1024
num_workers = 8
fake = false
data_world_size = 4

[ckpt]
path = "./"
interval = 18865

I would appreciate any guidance or insights you might have on resolving this issue. Specifically, any suggestions on further debugging steps or configuration adjustments that might help stabilize the training process would be invaluable. I'm quite new to this kind of setup, so any guidance would be greatly appreciated. Thank you in advance for your assistance!