Source-linked AI summary
A Scalable Pipeline for LLM-Teacher Distillation Labeling: Work-Stealing Job Scheduling and Memory-Aware GPU Concurrency
Ravi Satya Durga Prasad Yenugula
TL;DR
Large-scale LLM labeling needs both cost-quality evidence and GPU scheduling that remains effective under skew and worker failures. This paper combines relabel-gold benchmarking with work-stealing and memory-aware concurrency, achieving up to 3.43× static-sharding throughput under skew while completing all 2,000 tasks after half the workers were killed.
Problem
Large-scale LLM labeling lacks evidence about teacher quality per dollar and methods for keeping GPU workers utilized under skewed, failure-prone workloads.
Method
The paper combines relabel-gold benchmarking with a fault-tolerant work-stealing pool and device-aware concurrency sizing.
Results
3.43× throughput over static sharding at 0.9 skew, while work stealing completed all 2,000 tasks after half the workers were killed.
Takeaways & Limitations
The benchmark shows that a small teacher can be adequate for sentiment but inadequate for irony, while work stealing preserves throughput under severe skew and recovery after worker kills.
Takeaways & Limitations
Portability beyond the evaluated single-machine SQLite deployment is untested, including object-store and multi-node implementations.
Abstract
from arXiv · showhide
Labeling large text corpora with LLM teachers has become a practical route to training data at scale. At millions of items, hand-labeling every batch is not feasible, and two questions dominate: what label quality a teacher buys per dollar, and how to keep a fleet of GPU workers busy under skewed, failure-prone workloads. We present a simple, reproducible pipeline that addresses both. First, a work-stealing ring pool: each worker owns a queue, drains it first, and then steals from ring successors, with exactly-once task claims via atomic conditional writes and crash tolerance via stale-claim sweeping. The claim protocol requires only a compare-and-set primitive from its storage layer; we implement it on a single SQLite file, which makes the reference implementation dependency-free and the experiments reproducible on one machine. Second, a memory-aware concurrency rule that sizes per-node parallelism by how many model copies fit on the GPU, so the same code runs safely across device sizes. Third, a relabeling benchmark methodology in which the teacher relabels a public dataset that already has gold labels, so quality reduces to an agreement measurement and cost follows from measured throughput. Under skewed load the pool sustains up to 3.4 times the throughput of static sharding while matching it at zero skew, loses 0 of 2,000 tasks when half the workers are killed mid-run (static sharding loses 953), and yields measured quality and cost points for an instruction-tuned teacher on irony and sentiment tasks. All experiments run on public data and commodity hardware; code, tests, and run logs are released.
1 Introduction
The section frames scalable LLM labeling around teacher quality per compute cost and reliable GPU utilization under skewed, failure-prone workloads. It introduces a coordination-free pipeline combining work stealing, memory-aware concurrency, relabel-gold benchmarking, and a dependency-free implementation.
- Motivation: Millions of incoming items make hand-labeling infeasible, leaving teacher quality per compute dollar and GPU utilization under skew and preemption as central deployment questions.These conditions are described as typical of spot capacity.
- Pipeline design: Work stealing and conditional writes compose into a coordination-free labeling pipeline whose behavior can be measured by relabeling an already labeled dataset.The storage primitive enables exactly-once claims without a coordinator process.
- Contributions: The work-stealing ring pool provides atomic exactly-once claims and stale-claim sweeping for fault tolerance, using a single SQLite file and only a compare-and-set storage primitive.This is presented as the first systems contribution.
- Contributions: Memory-aware concurrency sizes per-node parallelism from device and model sizes, keeping the same code safe across 16 to 24 GB devices.The rule adapts parallelism to available memory rather than requiring separate code paths.
- Contributions: Relabel-gold benchmarking returns measured quality and cost points, identifying when a small instruction teacher is adequate versus when a larger teacher or human loop is required.The methodology is paired with a dependency-free reference implementation and committed run artifacts.
2 Related Work
Prior work covers work stealing, weak supervision and LLM annotation, and adapter-based fine-tuning, while this paper positions its contribution as combining these areas for LLM labeling. It emphasizes exactly-once coordination, relabel-gold evaluation, and measured cost-quality points rather than novelty in any individual method.
- Work stealing and distributed queues: Work stealing originates with Cilk, while ring stealing simplifies victim selection for OS-process workers through predictable contention and few-row SQL claims.Ring stealing prefers a fixed successor rather than a uniformly random victim.
- Weak supervision and LLM-as-annotator: Snorkel established programmatic labeling, while later studies measured prompted LLM agreement with human gold labels without addressing systems concerns.The cited systems lineage identifies Snorkel and its follow-ups as the closest labeling systems work.
- Adapters and quantized fine-tuning: QLoRA and DoRA enable adapter fine-tuning on commodity GPUs, supporting multiple teacher copies per device under the paper’s concurrency rule.The paper explicitly treats adapter methods as an implementation detail rather than a novelty claim.
- Positioning: The paper claims novelty in combining these components for LLM labeling, with exactly-once claims, relabel-gold cost-quality evaluation, and measured operating points.It disclaims novelty for work stealing, weak supervision, and adapters individually.
3 Method
The method combines an exactly-once work-stealing pool with stale-claim recovery, memory-aware teacher concurrency, and a relabel-gold benchmark that measures agreement and cost from observed throughput. It also defines deployment and failure-handling boundaries for SQLite-backed coordination and preemptible workers.
- Work-stealing pool: Workers drain their own queues before stealing from ring successors, using an atomic pending-to-running conditional write for exactly-once claims.The implementation uses SQLite IMMEDIATE transactions; the same conditional-write primitive can be provided by object stores.
- Failure recovery: Stale-claim sweeping resets abandoned running tasks to pending after a threshold, while late completions are discarded safely.The reported runs use a 60 s threshold; the mechanism is designed for workers disappearing from preemptible instances.
- Memory-aware concurrency: Concurrency is bounded by GPU memory for model weights and activation peaks and by available host CPU cores, using total rather than instantaneous free memory.With the same budget and code, a 16 GB device yields 2-way concurrency and a 24 GB device yields 4-way concurrency.
- Scope and deployment: The experiments measure the pool and single-teacher relabeling stage, not human evaluation or judge training, and the SQLite deployment assumes shared filesystem access.Object-store portability removes the filesystem constraint but increases conditional-write latency, which larger chunks can absorb.
- Relabel-gold benchmark: The relabel-gold benchmark compares teacher predictions on a public gold-labeled dataset using agreement and macro-F1, while converting measured throughput into dollars per 1,000 items.Tasks contain 50 texts so claim overhead remains small relative to inference, and concurrency follows the memory-aware rule.
- Throughput and cost: For chunk size B, coordination adds tc/B per item and larger chunks improve overhead but reduce load-balancing and recovery granularity.Measured cost per 1,000 items is 1000 P/(3600 R), where P is hourly price and R is aggregate throughput through the pool.
4 Experiments
Experiments show that work stealing preserves balanced throughput, sustains performance under skew, and recovers all tasks after worker failures. On the same commodity device, memory-aware concurrency safely runs four teacher copies while quality and cost vary sharply by task.
- Experimental setup: Throughput and fault-tolerance tests use CPU-bound synthetic work items, while quality and cost tests use flan-t5-base on a 24 GB NVIDIA A10G with 4 vCPUs.Each throughput measurement uses a fresh SQLite file and process pool; fault-tolerance runs kill half the workers after 300 ms.
- Throughput under skew: Work stealing matches static sharding at zero skew and remains essentially flat as skew increases, while static sharding degrades to the hot-queue rate.Experiments sweep W ∈{2, 4, 8} workers and skew from balanced routing to 0.9, where nine of ten tasks target Q0.
- Fault tolerance: 953 of 2,000 tasks are lost under static sharding after two of four workers are killed, whereas work stealing with stale-claim sweeping completes all 2,000.The sweep re-enqueues stranded tasks after Δ, allowing surviving workers to process them.
- Teacher quality and cost: 94.7% agreement with gold costs $0.0022 per 1,000 items for SST-2 sentiment, while tweet_eval irony reaches 49.6% agreement at chance.Both measurements use flan-t5-base through the pool with four teacher copies on the same 24 GB device.
- Memory-aware concurrency: 4 teacher copies run with zero outof-memory events on a 24 GB device under the memory-aware rule.Observed memory grows linearly at 0.58 GB per flan-t5-base copy; with a 2.0 GB per-copy budget, the rule predicts 4 safe copies after the core cap.
5 Discussion and Limitations
The pipeline’s limitations are concentrated in SQLite contention, unmeasured fleet-scale portability, task-dependent teacher quality, and an artificially concentrated skew model. These conditions constrain how broadly the reported performance and deployment claims should be generalized.
- Single-writer contention: SQLite serializes file-level writes, making the single-row claim update suitable when teacher inference dominates but potentially limiting faster-call or very large-fleet workloads.Observed claim time is small relative to teacher inference; an object-store backend is recommended when teacher calls take milliseconds rather than tens of milliseconds or fleets become very large.
- Portability is argued, not measured: The implementation and evaluation cover only SQLite on one machine, leaving the object-store port and multi-node deployment unshipped, unbenchmarked, and untested.The claim discipline is presented as portable by design, but fleet-scale portability remains a future-work claim.
- Task-dependent teacher quality: Teacher quality is task-dependent: the methodology measures agreement with gold labels rather than improving quality, with small teachers deployment-ready on some tasks but not others.The SST-2 versus irony comparison is intended to help practitioners assess suitability before labeling millions of items.
- Skew model: The skew experiment concentrates tasks on a single queue Q0, representing a worst case for static sharding and not the more diffuse skew seen in real deployments.Under diffuse skew, the static-versus-stealing gap is expected to be narrower, but this condition was not measured.
6 Reproducibility
The paper provides a compact, MIT-licensed reference implementation with offline validation, reproducibility scripts, and committed run artifacts. Its scheduler and fault-tolerance experiments can be regenerated within minutes on CPU.
- Implementation: MIT-licensed reference implementation is approximately 500 lines of Python.The implementation is designed to remain compact and dependency-light.
- Reproduction: Three scripts reproduce the benchmark, offline labeling, and memory-validation experiments.The commands target Sections 4.2, 4.3, and 4.5, while the fake teacher validates the full pipeline offline.
- Artifacts: All CSV and JSON run artifacts are committed under runs/ with timestamps matching the paper’s tables.This provides directly inspectable artifacts for comparing regenerated results with reported tables.
- Compute cost: Scheduler and fault-tolerance tables regenerate in minutes on CPU.The reported compute cost is described as modest.
A Full-precision throughput and fault-tolerance measurements
The throughput sweep reports throughput and speedup alongside underlying makespans and completion counts, with three repeats quantifying run-to-run variance. Variability is below 2% in half of the tested cells and reaches 10.7% in the worst case.
- Throughput variance: Three full-sweep repeats measured run-to-run throughput variance across 18 mode, W, and skew cells.Each Table 1 cell is a single run; the repeats used sample standard deviation with ddof=1 over the mean.
- Throughput variance: Under 2% throughput coefficient of variation occurred in 9 of 18 cells, while the worst case reached 10.7%.The coefficient of variation was computed across the three repeats.
B Choice of per-copy memory budget
The per-copy memory budget was set to 2.0 GB, about 3.4× the measured 0.58 GB weight footprint, based on profiling activation-memory behavior during generation.
- B Choice of per-copy memory budget: 2.0 GB corresponds to roughly 3.4× the measured 0.58 GB per-copy weight footprint.The multiplier was selected for the Section 4.5 budget.
- B Choice of per-copy memory budget: Concurrency was increased during single-worker profiling until per-batch activation memory during generate() plateaued.Profiling used flan-t5-base with 128-token inputs and 4 newly generated tokens per item.
- B Choice of per-copy memory budget: For flan-t5-base under that setup, activation memory does not exceed the weight footprint, while larger multipliers are safe but leave device memory unused.The stated setup used 128-token input and 4 new tokens generated per item.
C Notes on the SQLite backend
The SQLite backend uses WAL mode and a 30-second busy timeout, while maintaining low claim contention and a small, disposable database footprint during experiments.
- SQLite backend: WAL journal mode with a 30-second busy timeout supports the pool’s SQLite claim path.Under W=8 workers, the claim path issues approximately one BEGIN IMMEDIATE transaction per task.
- SQLite backend: Under W=8 workers, contention stays well under one lost-claim retry per hundred successful claims, requiring no explicit backoff.
- SQLite backend: The database file for a 2,000-task run occupies about 500 KB and is deleted at the end of each experiment.