A training cluster for a pharma research customer

Corvane Therapeutics (fictional), a pharma research company, wants to pretrain a 34-billion-parameter language model on 1.5 trillion tokens of chemistry and biomedical text within 30 days. They will buy servers with 8 H100 SXM GPUs each.

Assume:

• Per GPU: 989 TFLOPS of dense BF16 compute, 80 GB of memory at 3.35 TB/s, NVLink inside the server.
• Per server: 8 network ports at 400 Gb/s, one per GPU. Up to 10.2 kW at full load.
• Switches: 64 ports of 400 Gb/s each, available for InfiniBand or Ethernet.
• Training setup: sequences of 4,096 tokens, a global batch of about 4 million tokens, BF16 mixed precision with the Adam optimizer.
• A shared parallel file system gives this project 20 GB/s of write bandwidth.

Constraints:

• The 30-day window starts when the cluster passes acceptance. It cannot move, because the results feed a board decision.
• The data hall has 600 kW for this project. Network, storage and management take about 60 kW of it.
• Part of the text is licensed and must not leave the customer's site.
• The IT team knows Ethernet well and has never run InfiniBand.

Asks:

1. How many GPUs do they need? Show the math and your assumptions.
2. Propose a parallelism layout. Estimate the network traffic per GPU per training step, and how long it takes.
3. Design storage for the training data and for checkpoints.
4. Design the compute fabric: InfiniBand or RoCE, the topology, and the number of switches and links.

Write down your GPU count with the utilization you assumed, the network traffic per GPU per step, and the three risks most likely to make the customer miss the 30 days.