BEAM search 70B llama on tinybox

[chenyuxyz]

70B llama uses almost all vram and beam search new buffer allocation would fail due to out of resource.

It would be nice if we can search live with some flags to control buffer allocation. If not, we can have a regular run and log the kernels, and use a separate script to loop through the kernels and beam search them into cache.

[chenyuxyz]

see a weird error while trying to beam search llama 34B DEBUG=2 BEAM=2 PYTHONPATH=. python3 examples/llama.py --gen code --size '34B' --temperature 0 --shard 6 --count 10 --prompt "from tinygrad.tensor import Tensor" --timing

   3.08s:       377.80 us          7/  20           86  683    4    2   16    4HW Exception by GPU node-4 (Agent handle: 0x55d987988820) reason :GPU Hang
HW Exception by GPU node-4 (Agent handle: 0x559da9432ce0) reason :GPU Hang
HW Exception by GPU node-4 (Agent handle: 0x55d9a6070cb0) reason :GPU Hang
CLDevice: got 1 platforms and 0 devices
/usr/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 6 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
[1]    112769 IOT instruction (core dumped)  PARALLEL=1 DEBUG=2 BEAM=2 PYTHONPATH=. python3 examples/llama.py --gen code

that CLDevice: got 1 platforms and 0 devices means we are initializing a GPU device?