Issue 1:
We've built jaxlib 0.3.7 on conda-forge with cuda: conda-forge/jaxlib-feedstock#97

I am opening this issue for viz but also for testing/confirmation all is good. Could someone from the team here have a look with us? Easy things to check, we get something like ~11000 actions/cycles of bazel compilation (compared to ~7500 for non-cuda) and the size of our artifact is similar to that of the official jaxlib ~200 mb. I tested it locally and it seems to put the arrays on the GPU by default. So it is working, but more confirmation is always good!

Note that we build for 10.2, 11.0, 11.1, and 11.2 on conda-forge by default, though we we can adjust this if requested, and do note the forward compatibility of 11.2, it is really "11.2+" meaning that it compiles for 11.2 and above (if someone installs that the 11.2 compiled version, it will in fact install cudatoolkit 11.6 with it and the latest cudnn)

People interested in checking this out, should head to the:

PR: link above

logs: https://dev.azure.com/conda-forge/feedstock-builds/_build/results?buildId=506049

artifacts: https://dev.azure.com/conda-forge/feedstock-builds/_build/results?buildId=506049&view=artifacts&pathAsName=false&type=publishedArtifacts

You will see we still have a few problems (lost connections + timeouts mainly) due to the strict 6-hour time limit (it is tricky because we are very close to the limit, but hopefully we will manage to get all of them by rerunning the CI a few times)

Tagging @hawkinsp again who's been involved in issues related to the builds before + @jakevdp who's been very helpful and engaging.

Please:

Check for duplicate requests.

Describe your goal, and if possible provide a code snippet with a motivating example.

Issue 2:
Duplicate issues: #788 and #4528 are not suitable for this case.

How to reproduce the bug:

[1] Just run mu2Net on 8gpus A100, use BENCHMARK = 'ViT large / Chars benchmark'

[2] OOM error will occur when train_step function compiled by jax.jit is executed.

[3] The A100 have sufficent 80GiB memory.

Core code:

@partial(jax.jit, static_argnames=['model', 'optimizer'], donate_argnums=[0, 2])
def train_step(params, fixed_params, opt_state, images, labels, model, optimizer):
def loss_fn(params, fixed_params, images, labels):
logits = model.apply({'params': format_params(params, fixed_params)},
images, train=USE_DROPOUT)
labels = jax.nn.one_hot(labels, logits.shape[-1])
return -jnp.mean(jnp.sum(labels * nn.log_softmax(logits), axis=-1))
grads = jax.grad(loss_fn)(params, fixed_params, images, labels)
updates, opt_state = optimizer.update(grads, opt_state, params=params)
params = optax.apply_updates(params, updates)
return params, opt_state

def train_loop(paths, ds_train, ds_validation, devices, exp_config):
global LOOP_START
global TRAIN_PARAMS_TRANSFER
global TRAIN
timing = {'start_time': time.time(),
'start_time_loop': LOOP_START}
task = paths[0].task
# The following values should be shared by all paths in this generation batch.
for path in paths:
assert task == path.task
assert paths[0].hparams['ds_image_size'] == path.hparams['ds_image_size']

gc.collect()

# Compile.
compile_train_batches_arr = jax.device_put_replicated(
get_sample_batch(
paths[0].hparams['ds_image_size'],
task.train_batch_size),
devices)
compile_eval_batches_arr = jax.device_put_replicated(
get_sample_batch(
paths[0].hparams['ds_image_size'],
task.validation_batch_size),
devices)

for p_id, path in enumerate(paths):
if VERBOSE:
print('Parent')
print(prp(path.parent))
print(prp(path))
path.device_id = p_id % len(devices)
path.device = devices[path.device_id]
print("path:", p_id, "device:", path.device)
path.optimizer = path.get_optimizer()
path.optimizer_init_fn = jax.jit(path.optimizer.init, device=path.device)
path.best_params_local = None
path.best_opt_state_local = None
path.best_quality = None
path.best_score = path.parent.score() if path.task is path.parent.task else -np.inf
path.evals = []

# Launch parallel compilation of eval and train step functions.
params_local = path.get_trainable_params()
check_is_local(params_local)
path.compile_params_device = jax.device_put(params_local, path.device)
path.compile_fixed_params_device = jax.device_put(
path.get_fixed_params(),
path.device)
path.compile_train = Thread(
target=train_step,
args=(path.compile_params_device,
path.compile_fixed_params_device,
path.optimizer_init_fn(params_local),
compile_train_batches_arr['image'][path.device_id],
compile_train_batches_arr['label'][path.device_id],
path.model,
path.optimizer))
path.compile_eval = Thread(
target=eval_step,
args=(format_params(
path.compile_params_device,
path.compile_fixed_params_device),
compile_eval_batches_arr['image'][path.device_id],
compile_eval_batches_arr['label'][path.device_id],
path.model))
path.compile_eval.start()

for path in paths:
path.compile_eval.join()
del path.compile_eval
timing['end_compile_eval'] = time.time()
path.compile_train.start()
del compile_eval_batches_arr

for path in paths:
path.compile_train.join()
del path.compile_train
del path.compile_params_device
del path.compile_fixed_params_device
timing['end_compile'] = time.time()
del compile_train_batches_arr

gc.collect()

# Parameter transfer.
transfer_start = time.time()
for path in paths:
path.params_device = jax.device_put(
path.get_trainable_params(),
path.device)
path.fixed_params_device = jax.device_put(
path.get_fixed_params(),
path.device)
path.opt_state_device = path.optimizer_init_fn(path.params_device)
# Set opt state.
for c in path.components:
if c.is_trainable():
assert c.name in path.opt_state_device[1][0].trace.keys()
if c.opt_state is not None:
path.opt_state_device = (
path.opt_state_device[0],
(optax.TraceState(
trace=path.opt_state_device[1][0].trace.copy(
{c.name: jax.device_put(c.opt_state,
path.device)})),
path.opt_state_device[1][1]
)
)
check_is_on_device(path.opt_state_device, path.device)

TRAIN_PARAMS_TRANSFER += time.time() - transfer_start

iter_ds_validation = iter(ds_validation)
# TRAIN
for t_step, train_batch in zip(
range(exp_config.num_validations_per_path_training
* task.num_train_batches_between_validations),
ds_train,
):
replicated_start = time.time()
train_batch_arr = jax.device_put_replicated(train_batch, devices)
TRAIN_PARAMS_TRANSFER += time.time() - replicated_start
for p_id, path in enumerate(paths):
if t_step == 0:
timing['end_prep'] = time.time()
t_step_0_time = time.time()
train_step_start = time.time()
path.params_device, path.opt_state_device = train_step(
path.params_device,
path.fixed_params_device,
path.opt_state_device,
train_batch_arr['image'][path.device_id],
train_batch_arr['label'][path.device_id],
path.model,
path.optimizer)
TRAIN += time.time() - train_step_start
if t_step == 0 and time.time() - t_step_0_time > 1:
print(f'WARNING: First train step took: {time.time()-t_step_0_time:.2f} s')
del train_batch, train_batch_arr

# EVAL
# ...

Full error messages/tracebacks:

Exception in thread Thread-14:
Traceback (most recent call last):
File "/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/threading.py", line 890, in _bootstrap
self._bootstrap_inner()
File "/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/threading.py", line 926, in _bootstrap_inner
self.run()
File "/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/threading.py", line 870, in run
self._target(*self._args, **self._kwargs)
File "/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/api.py", line 476, in cache_miss
donated_invars=donated_invars, inline=inline, keep_unused=keep_unused)
File "/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/core.py", line 1765, in bind
return call_bind(self, fun, *args, **params)
File "/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/core.py", line 1781, in call_bind
outs = top_trace.process_call(primitive, fun_, tracers, params)
File "/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/core.py", line 678, in process_call
return primitive.impl(f, *tracers, **params)
File "/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/dispatch.py", line 185, in _xla_call_impl
return compiled_fun(*args)
File "/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/dispatch.py", line 615, in _execute_compiled
out_bufs_flat = compiled.execute(input_bufs_flat)
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 122875791936 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
parameter allocation: 1.43GiB
constant allocation: 8B
maybe_live_out allocation: 581.19MiB
preallocated temp allocation: 114.44GiB
preallocated temp fragmentation: 146.50MiB (0.13%)
total allocation: 115.86GiB
total fragmentation: 146.53MiB (0.12%)
Peak buffers:
Buffer 1:
Size: 1.27GiB
XLA Label: custom-call
Shape: f32[64,16,577,577]
==========================

Buffer 2:
Size: 1.27GiB
XLA Label: custom-call
Shape: f32[64,16,577,577]
==========================

Buffer 3:
Size: 1.27GiB
XLA Label: custom-call
Shape: f32[64,16,577,577]
==========================

Buffer 4:
Size: 1.27GiB
XLA Label: custom-call
Shape: f32[64,16,577,577]
==========================

Buffer 5:
Size: 1.27GiB
XLA Label: custom-call
Shape: f32[64,16,577,577]
==========================

Buffer 6:
Size: 1.27GiB
XLA Label: custom-call
Shape: f32[64,16,577,577]
==========================

environment:

python: 3.7.13

jax: 0.3.14

jaxlib: 0.3.14+cuda11.cudnn82