Training on a rented GPU: four and a half times cheaper

Every post so far has trained the fighter on Amazon EC2 spot machines with 96 or 192 vCPUs. Tony Yum, the coach, on where that was going:

As I’m provisioning more and more virtual CPUs on EC2 spot instances, the cost is racking up. And this is only the beginning. There are a lot more training to do, possibly bigger networks. So it’s only going to get more expensive. So let’s think of a way to save money.

GPU is one option, but there are some complications with the number of contacts and the way the body is built. So some experiments are needed. Yes, in order to save costs, I’ll have to throw even more money at an experiment, but it’ll be a good investment, I believe.

So this post isn’t about a new skill. It is about the machine the skills are learnt on: whether the training is cheaper on a rented GPU, what the code had to change to get there, and which GPU is cheapest. The models trained here were thrown away; only the timings and the bills mattered.

It covers one morning, 04:05 to 07:54 UTC on 21 October.

Why a GPU wasn’t the obvious answer

In the blocking post a rented GPU was tried for a few minutes and set aside: training waits on the physics, and the physics ran on the CPU. A GPU only helped the small network.

The way round that is to move the physics too. MuJoCo Warp is a version of MuJoCo that runs on the GPU and simulates thousands of copies of the world at once, all in step. That has two catches for this fighter.

• The body. The muscular body has 89 degrees of freedom in one tree, from the pelvis out to every finger joint, and 63 egg-shaped muscles on top. GPU simulators are built for robots with a few dozen joints. The simple body (Figure 1, right) has 43.
• The contacts. On a CPU each copy of the world keeps as many contacts as it needs. On the GPU every copy gets a fixed number of slots, set before training starts. After a fall, a muscular body lying on the floor reaches 32 contacts and 260 constraint rows, the equations the physics solves to keep things from passing through each other. Smaller slots overflowed and crashed the run; the slots that work leave most of the memory empty most of the time.

The test

The task was the first one the fighter ever learnt: stand in the stance, keep the rhythm, and stay up while being shoved (the balance post). Each machine trained it from nothing until it was solved: amplitude at full, at most 5 % of episodes ending in a fall, and episodes lasting the full 10 s. Both bodies were tested, muscular-fast and simple.

Two figures were kept for every machine: time to solve, and cost per solve, which is that time times the machine’s price per hour.

Both sides train many copies of the world at once, parallel environments. The CPU side was two ordinary EC2 spot runs with the usual trainer: 72 worker processes with 6 copies each.

machine	body	price	time to solve	cost
c8i.24xlarge spot, 96 vCPU	muscular-fast	$1.45/h	714 s	$0.29
c8a.24xlarge spot, 96 vCPU	simple	$2.02/h	404 s	$0.23

Step 1: tuning the code for one GPU

The first GPU was an RTX 5090, the newest consumer card, rented on vast.ai for $0.74 an hour. vast.ai is a market where people and small data centres rent out their GPU machines by the second.

The GPU trainer was written two days earlier and tested on an older card. On the 5090 it ran 44,000–47,000 steps a second with 4,096 copies of the world, the GPU at 97 %. The physics took nine tenths of the time.

Finding the slow part. Profiling timed every small program the GPU ran (each one a kernel). One took 53 % of the time: the factorisation of the mass matrix, the table that says how hard each joint is to move given everything attached to it. MuJoCo Warp has a fast method for bodies up to 64 degrees of freedom. The muscular body’s 89, all in one tree, fell back to a slow general method: 24 separate launches per factorisation, one for each level of the tree, twice in every physics step.

A replacement. The new version (warp_factor.py) gives each copy of the world a group of 64 GPU threads that do the whole factorisation in one launch, waiting for each other between levels instead of starting a new kernel. Long sums are split into chunks of 8 so no thread is left doing all the work. The factorisation went from 38.6 to 8.9 ms of GPU time per step, and the physics got 28 % faster: 76,000 steps a second at 8,192 copies. It is the same physics: on the same episode with the same actions, the joint angles matched the stock method to a millionth of a radian, as close as 32-bit numbers allow.

The simple body didn’t need this. Its 43 degrees of freedom fit MuJoCo Warp’s fast method.

Then the surprise: fewer copies of the world trained faster. With the new factorisation and 8,192 copies, the GPU trained at 61,000 steps a second, five times the CPU machine. After 12 minutes the fighter had reached only 60 % of the full stance. The CPU machine had solved it in less.

The reason is the curriculum. Each time the stance gets wider, it waits a fixed number of training iterations, about two episodes for every copy of the world, before it can rise again. Each iteration steps every copy 32 times. So a solve takes about 350–500 iterations whatever the machine, and the time to solve is that number times the seconds per iteration. More copies only make each iteration take longer. Training is fastest with as few copies as keep the GPU busy, and on the 5090 that was 1,024 (Figure 2).

At 1,024 copies the GPU is no longer short of work to do; it is waiting. Each of the thousands of small kernels in a physics step has a fixed start-up cost, and with 1,024 copies that cost sets the pace, not the arithmetic. So the next change cut the waiting.

No questions to the CPU in the middle of a step. Code like “did any copy fall?” makes the CPU stop and wait for the GPU’s answer. Falls, resets, shoves and foot forces were rewritten to work on every copy at once with a mask, and the trainer reads its figures from the GPU once an iteration.

Muscular-fast at 1,024 copies stayed at 31,000–32,000 steps a second on the 5090. On that machine the physics was the limit. The waiting mattered on the next machines.

Step 2: nine GPUs

A script (vast_bench.py) rented the cheapest offer of each GPU type at the same time, copied the code over, ran the same speed tests and two minutes of training per body, and handed the machine back. Nine types, from a 2018 data-centre T4 at $0.15 an hour to an RTX PRO 6000 at $1.00.

The first pass was a disappointment. Every machine except the 5090’s took 0.7–1.4 s per iteration to update the network, where the 5090 took 0.2 s. The 5090 had happened to sit in a desktop with a fast new processor. Most rented GPUs sit in servers with slower, older processors, and the trainer was waiting on them. An L40S, a card nearly as fast as the 5090 here, trained at 13,000 steps a second.

So the CPU was taken out of the loop:

• The network update as a CUDA graph. One training step on a minibatch (run the network, work out the gradients, clip them, update the weights) is recorded once and replayed as a single unit, instead of the CPU sending each of its hundreds of kernels one at a time. 0.6–0.7 s an iteration became 0.2 s.
• The whole environment step as a CUDA graph, the physics included. Only 5 % on its own.
• One lookup instead of eight. The reference motion (where each joint should be in the stance’s rhythm) was 8 tables looked up one by one, three times a step: 8 of the 21 ms the step spent outside the physics. Now the tables sit side by side and one lookup reads them all, and the non-physics work fell to 8.7 ms. With that and the contact slots halved for the simple body (the foot-force check reads every slot), the simple body went from 41,000 to 48,000 steps a second on an RTX 3090.

Two bugs turned up on the way. The first step of every run had been simulated twice, once to prepare the GPU code and once for real (harmless, now fixed). And once, a heap of fallen muscular bodies ran out of constraint slots, the second catch above. The slots went up from 320 to 512 rows, at no measurable cost, and the trainer now prints OVERFLOW when any copy runs out.

With the new code, a second pass of the survey (Figure 3) gave cards behind ordinary server processors their real speed. A second test on the two cheapest, on different machines in other countries, gave the same seconds per iteration within a few percent.

Step 3: GPU against EC2

The full list, cost per solve with set-up included. ”~” marks an estimated time.

machine	where	$/h	muscular-fast	$	simple	$
RTX 5080	vast.ai	0.39	515 s	0.065	230 s	0.034
L40S	vast.ai	0.44	452 s	0.065	227 s	0.038
RTX 3090	vast.ai	0.20	1163 s	0.070	451 s	0.030
RTX 4090	vast.ai	0.51	436 s	0.074	217 s	0.043
RTX 5090 (fastest)	vast.ai	0.69	377 s	0.089	185 s	0.053
T4	vast.ai	0.15	2326 s	0.100	1025 s	0.046
L4	vast.ai	0.36	1074 s	0.117	389 s	0.048
A100	vast.ai	0.46	931 s	0.129	403 s	0.062
g6.xlarge (L4)	EC2 spot	0.48	~1074 s	~0.167	~389 s	~0.076
g6e.xlarge (L40S)	EC2 spot	1.47	~452 s	~0.259	~227 s	~0.167
c8a.24xlarge, 96 vCPU	EC2 spot	1.34	714 s	0.299	404 s	0.184
c8a.48xlarge, 192 vCPU	EC2 spot	3.65	~529 s	~0.627	~299 s	~0.394

What it shows:

• A rented consumer GPU solves the stance for less than a quarter of the price: $0.065 on an RTX 5080 or L40S against $0.30 on the 96-vCPU spot machine for muscular-fast, $0.03–0.04 against $0.18 for simple. And it is faster too: the 5080, L40S and 4090 take 55–70 % of the CPU machine’s time, the 5090 about half.
• The cheapest cards are all within about 15 % of each other: RTX 5080, L40S, RTX 3090 and RTX 4090. The 3090 is the cheapest by the hour but more than twice as slow. The 5090 is the fastest at about 30 % more per solve, worth it when time matters more than money.
• Data-centre cards are poor value here. The A100, L4 and T4 are built for large blocks of matrix maths and lots of memory, but this physics is ordinary 32-bit arithmetic in thousands of small steps. It wants clock speed, not memory or tensor cores.
• EC2’s own GPU machines don’t pay. The same L40S costs 3.4 times as much on EC2 as on vast.ai. The best of them, g6.xlarge with an L4, costs about half the CPU machine per solve but is slower than it.
• The simple body costs about half as much on a GPU: it runs 1.7 times as many steps a second and needs about the same number of iterations. On the CPU the gap is only 1.3 times.

Getting a machine ready

Speed per iteration isn’t the whole bill. A rented machine is paid for from the moment it starts, before any training. On vast.ai the time spent downloading the software image is free; copying the code over, installing the physics libraries and preparing the GPU code is paid at the full rate.

At first the copy was slow. The script sent every file in the project, 57 MB, 51 MB of it a table only the walking environment uses, and many cheap hosts have slow incoming connections: the A100 in California took 17 minutes to receive it, then 26 minutes to install the libraries over the same line. Now the script sends 1.4 MB and asks for hosts with at least 500 Mbit/s. On a good host the paid set-up is about 1.5 minutes, one or two cents, about what an EC2 CPU machine takes from launch to the first iteration (75–93 s).

The real risk is a bad host. An RTX 4090 in Thailand ran a tenth as fast as another 4090 on both passes. An RTX PRO 6000 started, never answered, and cost $0.51 for nothing. An RTX 5080 had a driver too old for the GPU code and couldn’t see its own card. The script now asks for a recent driver and gives up on any machine that doesn’t answer within 5 minutes.

What Coach Tony thinks

Tony’s brief for the morning was to optimise the GPU code for the machine first, then compare: both bodies, other GPU machines to find the cheapest, and EC2. He called the experiment money spent to save money, “a good investment”. The models could be thrown away.

When the first survey came back, he asked why the uploads were so slow and wanted the set-up costs in the comparison, not just the training time. That is where the 57 MB upload and the bad-host losses came to light.

What’s next

Only the stance runs on the GPU so far. Punching, walking, blocking and the opponent environment each need their own GPU version before they can move off EC2. On these numbers that work pays for itself quickly: the last post’s 20 blocking runs cost about $33 on EC2.

What it cost

Most of the cost was renting GPUs to measure them, plus two CPU runs on EC2 to compare against, each left running a few minutes past its solve. The vast.ai figure comes from vast.ai’s charge records; the EC2 figure is an estimate from each run’s hours and spot price.

item	machine-hours	cost
RTX 5090 on vast.ai: profiling, the new factorisation, the copies-of-the-world tests, the simple body	about 1.7	about $1.21
RTX 3090 on vast.ai: the CPU-out-of-the-loop work and two runs from nothing	about 1.1	about $0.22
the survey: about 17 GPU machines over two passes and a timed re-run	about 5.5 (estimate)	about $2.61
2 EC2 spot runs (c8i.24xlarge and c8a.24xlarge, 96 vCPU)	about 0.8	about $1.38
total	about 9	about $5.40