[fix] small improvements for dataloader, metrics, optimizer robustness, and ops - #797
[fix] small improvements for dataloader, metrics, optimizer robustness, and ops#797sherlockwu wants to merge 1 commit into
Conversation
409cd25 to
a076eca
Compare
sherlockwu
changed the title
6cb1549 to
0e58308
Compare
sherlockwu
changed the title
sherlockwu
changed the title
0e58308 to
7a242d9
Compare
…ution metrics Infra-only changes, independent of any drafter architecture: - feature_dataloader: make dataloader_num_workers actually multi-threaded. Previously it only sized the buffer of finished batches — one background thread fetched a batch at a time, capping delivery at 1/fetch-latency regardless of buffer depth. The prefetch worker now submits fetches to a num_workers-sized pool and the buffer carries futures (measured on the 32K DSV4-Flash recipe: 122.6 s/step single-fetcher -> 31.4 s at 16) - optimizer: skip the step deterministically on a non-finite global grad norm instead of NaN-ing the weights permanently - controller + loader perf counters: perf/data_wait_producer_s, perf/data_wait_fetch_s, perf/fetch_seconds_per_sample, perf/fetch_delivered_gib_per_s — attribute trainer data waits to "capture cannot keep up" vs "transfer is slow" - dflash_family_model: masked_fill instead of a per-microbatch pageable H2D copy; device-side assert instead of a host sync after all_reduce Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
7a242d9 to
68e1820
Compare
sherlockwu
changed the title
sherlockwu
changed the title
…ttribution) into coloc-perf-base
| state.mark_materialized(refs) | ||
| if not put_interruptibly((batch, refs)): | ||
| future = pool.submit(self._timed_make_batch, refs) | ||
| state.track_fetch(future) |
There was a problem hiding this comment.
track_fetch() removes the future as soon as it completes, while mark_materialized() happens only when the consumer calls future.result().
So it would trigger a race?
| self._store_get_tensor(key, dst) | ||
| return dst | ||
| except KeyError as exc: | ||
| self._quarantined_buffers.append(dst) |
There was a problem hiding this comment.
This _quarantined_buffers will not be released?
| logger.warning("%s; retry %d/3 in %.0fs", exc, attempt + 1, delay) | ||
| time.sleep(delay) | ||
| dst = _alloc_from_spec(spec) | ||
| self._store_get_tensor(key, dst) |
There was a problem hiding this comment.
Why there is no more try catch?
| try: | ||
| self._store_get_tensor(key, dst) | ||
| return dst | ||
| except KeyError as exc: |
There was a problem hiding this comment.
Should this retry only the explicit transient TRANSFER_FAIL/-800 status?
| for mp in self.fp32_params: | ||
| mp.grad = None | ||
| self.last_grad_norm = grad_norm.detach() | ||
| return self.last_grad_norm |
There was a problem hiding this comment.
(nit)Here we still silently return, so there will be a empty step.
| "perf/data_wait_fetch_s": ( | ||
| loader["wait_fetch_s"] / steps | ||
| ), | ||
| "perf/fetch_seconds_per_sample": ( |
There was a problem hiding this comment.
This is per batch not per sample right?
As the denominator is batches
|
Measured this branch (merged onto main
So the pool converts the wait into overlapped fetches, but on one node the same ~2 GB/s reaches the trainers either way; the gain shows up when the consumer pulls from another node. Two operational notes from the same runs: with 8 fetches in flight the default 500 ms and a 3 s Mooncake lease both produced |
2 participants
Robustness
optimizer.py: a non-finite global grad norm now skips the optimizer step deterministically on every rank (clipping cannot rescale a NaN norm; one such Adam step NaNs the weights permanently).mooncake_store.py: transientget_intofailures (TRANSFER_FAIL/-800 under source-side contention) are retried with backoff instead of killing the run as aKeyError; the receive buffer of a failed attempt is quarantined — a timed-out transfer can still write into its destination after the call returns, and reusing that storage has corrupted the heap. Observed rate: ~9 per 100 steps on long 32K runs; an 8K validation run on this branch without the retry died within 12 steps.Performance
feature_dataloader.py: makedataloader_num_workersactually multi-threaded. Previously the value only sized the buffer of finished batches — a single background thread fetched one batch at a time, capping delivery at 1/fetch-latency no matter how deep the buffer. Now the prefetch worker submits fetches to anum_workers-sized thread pool and the buffer carries futures, so up tonum_workersbatches materialize concurrently while leasing continues.Useful for long context feature passing. On the 32K DSV4-Flash disaggregated recipe (2-rank consumer, 3× TP2 capture, identical stack; steady-state tails):
num_workers: 16dflash_family_model.py:masked_fillinstead of a per-microbatch pageable H2D copy; a device-side assert instead of a host sync right after a collective.Observability
controller.py+ loader perf counters:perf/data_wait_producer_s,perf/data_wait_fetch_s,perf/fetch_seconds_per_sample,perf/fetch_delivered_gib_per_s— attributes trainer data waits to "capture cannot keep up" vs "transfer is slow".