Spaces:
Paused
Paused
File size: 2,240 Bytes
2f86d9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | # Copyright 2026 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Safe metric writer that avoids NCCL crashes on multi-GPU.
Overrides ``write_param_overview`` to count parameters without
calling ``jax.device_get()`` on sharded arrays.
"""
import dataclasses
from kauldron.train import metric_writer
# --- Patch B: Safe _convert_leaf using single-shard reads ---
# Never call jax.device_get() on sharded arrays. Always read from a
# single local shard to avoid triggering NCCL AllGather.
import jax
import numpy as np
import kauldron.train.auxiliaries as aux
def _convert_leaf(leaf):
if isinstance(leaf, jax.Array):
shard_data = leaf.addressable_shards[0].data
return np.asarray(shard_data)
return leaf
aux._convert_leaf = _convert_leaf
# --- Patch C: Safe _compute_metric with synchronization ---
# Force jax.block_until_ready after metric computation to prevent
# concurrent NCCL operations that corrupt the communicator.
_orig_compute = aux.AuxiliariesState.compute
def _safe_compute(self, *args, **kwargs):
result = _orig_compute(self, *args, **kwargs)
jax.block_until_ready(jax.tree_util.tree_leaves(result))
return result
aux.AuxiliariesState.compute = _safe_compute
@dataclasses.dataclass(frozen=True, eq=True, kw_only=True)
class SafeMetricWriter(metric_writer.KDMetricWriter):
"""Metric writer that avoids NCCL crashes on multi-GPU.
Overrides ``write_param_overview`` to count parameters without
calling ``jax.device_get()`` on sharded arrays.
"""
def write_param_overview(self, step: int, params) -> None:
num_parameters = metric_writer.parameter_overview.count_parameters(params)
self.write_summaries(step=step, values={'num_params': num_parameters})
|