TACO: Efficient Communication Compression of Intermediate Tensors for Scalable Tensor-Parallel LLM Training
Abstract: Handling communication overhead in large-scale tensor-parallel training remains a critical challenge due to the dense, near-zero distributions of intermediate tensors, which exacerbate errors under frequent communication and introduce significant computational overhead during compression. To this end, we propose TACO (Tensor-parallel Adaptive COmmunication compression), a robust FP8-based framework for compressing TP intermediate tensors. First, we employ a data-driven reshaping strategy combined with an Adaptive Scale-Hadamard Transform to enable high-fidelity FP8 quantization, while its Dual-Scale Quantization mechanism ensures numerical stability throughout training. Second, we design a highly fused compression operator to reduce memory traffic and kernel launch overhead, allowing efficient overlap with communication. Finally, we integrate TACO with existing state-of-the-art methods for Data and Pipeline Parallelism to develop a compression-enabled 3D-parallel training framework. Detailed experiments on GPT models and Qwen model demonstrate up to 1.87X end-to-end throughput improvement while maintaining near-lossless accuracy, validating the effectiveness and efficiency of TACO in large-scale training.
Summary
- The paper introduces a novel adaptive scale–Hadamard transform combined with dual-scale quantization to effectively compress tensor-parallel communication.
- It demonstrates near-lossless convergence with less than 0.3% validation loss degradation on models ranging from 6.7B to 7B parameters.
- Fused GPU kernel implementation yields up to 1.87× throughput improvement in 3D-parallel training, addressing critical TP communication bottlenecks.
TACO: Communication Compression for Tensor-Parallel LLMs
Introduction and Motivation
Scaling LLMs to the trillion-parameter regime has made 3D parallel training—data parallel (DP), pipeline parallel (PP), and tensor parallel (TP)—the dominant methodology for feasible training. Within this paradigm, TP emerges as the primary performance limiter due to the heavy intra-node communication required to synchronize intermediate tensors during both forward and backward passes (Figure 1). Unlike DP (sparse, infrequent gradients) and PP (structured, low-synchronization activations), TP mandates tightly coupled, high-frequency exchange of dense activations and gradients. In recent benchmarks, more than half of the wall-clock training time in state-of-the-art systems is attributed to TP communication. Reducing this bottleneck is essential for further scaling.
Figure 1: Execution time breakdown across different models and scales, showing tensor parallel communication dominates as model size and parallelism increase.
The paper "TACO: Efficient Communication Compression of Intermediate Tensors for Scalable Tensor-Parallel LLM Training" (2604.24088) addresses the acute challenge of compressing these critical-path TP communications during training, which is distinct from both DP/PP and inference-oriented quantization settings. Unlike activations or gradients in other contexts, TP intermediate tensors have distributions that are sharply concentrated near zero, but feature long tails—rendering existing INT8 and even naïve FP8 approaches susceptible to catastrophic instability and divergence during extended training.
Analysis of TP Intermediate Tensor Distributions
Quantitative distributional analysis reveals that TP tensors deviate significantly from assumptions made in conventional quantization: most of the tensor mass resides in extreme proximity to zero, with a small fraction forming a heavy tail (Figure 2, Figure 3). INT8, with its uniform quantization levels, lacks the resolution to represent small-magnitude values without collision—resulting in high mean-squared-error and loss of crucial information for optimization (Figure 4). In contrast, exponent-based formats such as FP8 (E4M3/E5M2) provide denser point placement near zero. However, even off-the-shelf FP8 fails under repeated, high-frequency synchronization as it does not adaptively condition the underlying tensor statistics, especially during blockwise distributed training (Figure 5, Figure 4).
Figure 2: Histogram illustrates the highly skewed, zero-centered distribution of TP communication data.
Figure 3: Data distribution characteristics of INT8 (uniform) and FP8 (non-uniform, higher density near zero).
Figure 4: Quantization errors of INT8 are uniform but considerable near zero; FP8 reduces error, especially for small values relevant in TP tensors.
The TACO Framework
To resolve these limitations, TACO introduces a threefold system:
- Adaptive Scale–Hadamard (ASH) Transform: This step performs energy-equalizing rescaling at the block level (with block size typically set to 256), followed by an orthogonal Hadamard rotation. The result is to ‘flatten’ the block distribution, amplifying low-magnitude values so they become representable within the finite precision of the target quantization scheme (Figure 6). Use of RMS (rather than variance) provides a shift-invariant scaling. This transforms formerly problematic dense regions (collapsed by naïve quantizers) into a more FP8-friendly, Gaussian-like spread.
- Dual-Scale (DS) Quantization: After the blockwise ASH transformation, two scales are maintained per block: one for the adaptive pre-rotation rescaling, and another for post-rotation normalization such that all values fit within the FP8 format range.
- Fused System Kernels: Compression, scale computation, Hadamard transform, FP8 quantization, and communication scheduling are combined into a single, block-local GPU kernel (Figure 7, Figure 8), minimizing launch and memory traffic overhead.
Figure 8: TACO system overview; intermediate tensors are compressed with ASH and DS, packed with metadata, and decompressed using highly fused kernels.
Figure 6: Distribution of TP tensors before and after Hadamard-based transformations. AS-Hadamard uniquely redistributes zero-centered values, maximizing FP8’s dynamic range utilization.
Figure 7: Comparison of traditional multi-kernel execution (top) with TACO fused execution (bottom) — kernel fusion reduces memory traffic and increases efficiency.
Impact on Training Accuracy and Throughput
Under rigorous large-scale training, TACO achieves near-lossless convergence on GPT-6.7B and Qwen2.5-7B, with validation loss degradation of less than 0.3% relative to uncompressed baselines (Table 1, Figure 9). Component ablations show that ASH+DS is necessary: neither off-the-shelf NVFP8 nor Hadamard transform alone suffices (Figure 10, Figure 11).
INT8 quantization—even when combined with ASH conditioning—quickly results in catastrophic divergence, as the dynamic range remains insufficient for reshaped blocks (Figure 12). FP8 E4M3 is empirically preferred over E5M2: the former’s broader dynamic range at moderate precision better accommodates blockwise activation statistics encountered in TP.
Figure 10: Ablation: Stepwise introduction of ASH and DS progressively restores convergence, with full TACO nearly identical to the uncompressed baseline.
Figure 9: Validation loss for TACO and baseline on GPT 6.7B, Qwen2.5-7B, showing negligible gap.
Figure 11: Direct comparison of baseline, standard Hadamard, and ASH transforms — only ASH recovers near-baseline accuracy.
Figure 12: Validation loss curves for ASH+INT8, ASH+FP8(E5M2), and ASH+FP8(E4M3). INT8 diverges; E4M3 maintains stability and accuracy.
System Throughput and Scaling
TACO delivers up to 1.87× throughput improvement over uncompressed baselines under high TP degrees (Figure 13, Figure 14, Table "3D TFLOPS"). Kernel fusion is essential to these gains; naïve implementations without fusion lose much of the benefit (Figure 14).
As model and parallel degree increase—especially in full 3D-parallel stacks (TP+PP+DP)—enabling TP compression becomes disproportionately beneficial. When integrated with state-of-the-art DP and PP communication compressors (SDP4Bit, TahQuant), system-level results show that activating TACO in TP provides the bulk of the incremental performance (Table "3D TFLOPS", Figure 15).
Figure 13: End-to-end throughput (TFLOPS) under different TP degrees for GPT-2.7B/6.7B. TACO scales significantly better as parallelism intensity increases.
Figure 14: Throughput comparison; bar labels indicate speedup ratio relative to previous baselines.
Figure 15: Validation loss for GPT-6.7B under full 3D parallelism, showing TACO’s training is robust and stable.
Theoretical and Practical Implications
TACO makes several explicit, experimentally justified claims:
- INT8 is fundamentally incompatible with stable TP tensor compression during training due to inability to preserve dense near-zero information, even with blockwise adaptive transforms.
- FP8 (E4M3) with blockwise adaptive transform enables lossless-to-nearly-lossless TP communication compression up to at least billions of parameters, retaining training stability.
- System throughput is bottlenecked by TP in 3D parallelism; applying compression here yields super-linear speedup relative to compression in DP or PP.
This work offers practical guidance for both system architects and algorithm designers. It directly supports adoption of FP8 as the low-precision communication format in TP, contingent on careful statistical conditioning via blockwise adaptive transforms and dual-scaling.
Future Directions
The TACO methodology is, by construction, generalizable to new hardware generations, as the explicit transformations can be mapped to INT8 in the absence of native FP8 support (though at lower precision). Extension to further compressing gradients and optimizer state is a clear next step, as is exploration of layer- or time-variant adaptive quantization. Additionally, future work can target dynamic block size selection and more aggressive kernel fusion to further improve arithmetic intensity.
Conclusion
TACO provides a robust, data-driven solution to one of the most significant scalability challenges in modern LLM training: efficient, stable compression of TP intermediate tensors. By integrating adaptive blockwise reshaping (ASH), dual-scaling quantization, and highly optimized GPU kernels, TACO achieves significant speedup (up to 1.87×) with minimal accuracy loss in end-to-end 3D parallel LLM training. This work establishes new baselines for communication-efficient distributed training of billion-scale models and paves the way for robust, low-bit TP communication in next-generation systems.
Reference:
"TACO: Efficient Communication Compression of Intermediate Tensors for Scalable Tensor-Parallel LLM Training" (2604.24088)
Paper Prompts
Sign up for free to create and run prompts on this paper using GPT-5.
Top Community Prompts
Explain it Like I'm 14
Overview: What is this paper about?
This paper introduces TACO, a new way to speed up training very LLMs that run across many GPUs at once. When these huge models train, the GPUs have to constantly share “intermediate” data with each other. That sharing (communication) can take a lot of time and slow everything down. TACO compresses that data so it travels faster, without hurting the model’s learning. The authors show it can make training much quicker while keeping accuracy almost the same.
What questions were the researchers trying to answer?
- How can we reduce the time spent sending data between GPUs during training of giant LLMs?
- Can we compress these “in-between” tensors safely, even though they’re exchanged very often and tiny mistakes can build up?
- Is there a practical compression method that’s both fast on GPUs and stable enough to keep training from breaking?
How did they approach the problem?
The problem in simple terms
- Imagine a group project where teammates constantly pass notes with partial work to each other. If passing notes takes too long, the whole project slows down.
- In LLM training, this “note passing” happens in something called tensor parallelism (TP). GPUs split a layer’s work and must sync partial results every layer, forward and backward. This happens a lot, so it’s on the critical path.
- Many of these numbers are very close to zero. If you compress them the wrong way, you lose subtle but important details. Because this happens many times per step, tiny errors can pile up and make training unstable.
Key ideas in TACO
- Uses FP8 instead of INT8 for compression:
- INT8 treats all numbers with equal spacing (like marking every centimeter on a ruler). That’s bad when most numbers are very close to zero, because you lose fine detail there.
- FP8 is a tiny floating-point format. It keeps an exponent, so it can represent small numbers more precisely near zero, which fits these tensors better.
- Adaptive Scale–Hadamard (ASH) transform (reshapes the data before compression):
- Think of the data in small blocks. TACO first turns up or down a “volume knob” per block so quiet blocks aren’t lost and loud blocks don't clip.
- Then it applies a Hadamard transform (a fast, mix-and-spread operation) to spread values out more evenly, so FP8 can represent them better.
- This combo helps prevent “zero-collapse,” where many tiny values would otherwise round to zero.
- Dual-Scale Quantization (two protective knobs):
- TACO uses two scales per block:
- One scale (pre-rotation) to normalize the block’s energy so tiny values get amplified and aren’t lost.
- One scale (post-rotation) to make sure nothing overflows the FP8 range (avoiding “clipping”).
- On the receiving side, it reverses both steps to recover the original data.
- Fused, fast GPU kernels and overlap with communication:
- Instead of doing many separate steps (each reading/writing memory), TACO fuses them into a single, well-optimized GPU kernel to cut overhead.
- It overlaps compression work with data transfer, so GPUs are busy computing while also sending/receiving data.
- Works with full 3D parallel training:
- The authors integrate TACO with other state-of-the-art compression for data parallelism (DP) and pipeline parallelism (PP), so the whole training stack benefits.
What did they find and why is it important?
- On GPT-style models and Qwen2.5-7B, TACO:
- Improved end-to-end throughput for TP by up to 1.87×.
- Improved throughput in full 3D-parallel setups by up to 1.53×.
- Kept training “near-lossless,” meaning accuracy and convergence stayed stable.
- Why it matters:
- For very large models, communication can take over half of the training time. Cutting that significantly speeds up training and lowers costs.
- Prior compression methods worked better for DP or PP, but TP is trickier because it communicates very frequently. TACO shows how to compress TP safely during training, not just inference.
What’s the bigger impact?
- Faster, cheaper training of huge LLMs:
- If communication becomes less of a bottleneck, we can train larger models or train the same models faster.
- Practical and future-ready:
- FP8 is increasingly supported on modern GPUs, so this method fits today’s and tomorrow’s hardware.
- More stable low-precision training:
- The combination of ASH (reshaping the data) and Dual-Scale Quantization (two safety knobs) is a general recipe for stable compression where tiny errors would otherwise accumulate.
In short, TACO shows a smart, hardware-friendly way to compress the most communication-heavy part of LLM training, making it both faster and reliable.
Knowledge Gaps
Knowledge gaps, limitations, and open questions
The paper presents TACO for FP8-based compression of tensor-parallel (TP) intermediate tensors with promising results, but several aspects remain underexplored or uncertain:
- Scalability beyond intra-node and small-scale settings: no evidence for stability or speedups on large multi-node clusters (e.g., 64–1024 GPUs), diverse topologies (ring/tree/dragonfly), or congested networks; integration with in-network reduction (e.g., SHARP) is unaddressed.
- Hardware generalization: evaluation limited to NVIDIA H100 with native FP8; portability and performance on AMD MI300, older NVIDIA generations without FP8 (A100), PCIe-only interconnects, or mixed-vendor clusters are unknown.
- FP8 format policy: criteria for choosing E4M3 vs E5M2 (and potential dynamic switching) per-tensor, per-pass (forward vs backward), or per-layer are unspecified; no sensitivity analysis across formats.
- Long-horizon and large-scale training stability: convergence and accuracy under full pretraining for >10B–70B+ parameter models, longer token budgets, and diverse corpora are not demonstrated.
- Theoretical guarantees: no formal error bounds on ASH+Dual-Scale quantization through TP collectives, nor convergence guarantees under stochastic optimization; effect of per-block scaling on bias and optimizer dynamics (Adam/LAMB) remains uncharacterized.
- Hyperparameter sensitivity and autotuning: impact of block size B, target energy τ, ε, FP8 variant, overlap chunk size, and two-shot parameters on accuracy and throughput lacks systematic study or autotuning strategy.
- Partitioning strategy: only contiguous flattened blocks are used; whether semantically aligned partitions (per head/channel/sequence position) reduce error or improve stability remains unexplored.
- FWHT constraints: handling of non-power-of-two block sizes, padding-induced overhead, and accuracy effects are not discussed; alternatives (DCT/Haar) for arbitrary sizes are not evaluated.
- Metadata overhead: two scalars per block are assumed negligible; for small tensors or micro-batches, the metadata-to-payload ratio, its compressibility, and possible shared or quantized metadata strategies need quantification.
- Coverage of TP collectives: behavior under sequence parallelism (AllGather/ReduceScatter) and attention-specific pathways (e.g., K/V communication) is not fully evaluated; guidance on where TACO is most effective is missing.
- 3D-parallel interaction: cumulative error from combining TACO (TP) with SDP4bit (DP) and TahQuant (PP) lacks thorough analysis; joint scheduling/policy to allocate compression budgets across dimensions is absent.
- Training regimes beyond pretraining: robustness in finetuning, instruction tuning, and RLHF—where gradients are higher variance and datasets smaller—remains untested.
- Numerically sensitive submodules: impact on softmax attention stability, LayerNorm/RMSNorm, SwiGLU/GLU gating, and residual pathways is not quantified; layer-wise failure modes are unknown.
- MoE and expert-parallel workloads: applicability to AllToAll-heavy expert routing and expert gradients, including error accumulation across sparse pathways, is not studied.
- Power and cost efficiency: net energy savings considering added compute (FWHT, reductions, conversions) versus reduced communication are not reported; roofline or energy-per-sample analysis is missing.
- NCCL/COCCL integration and maintainability: reliance on COCCL raises questions about availability, stability, and compatibility with mainstream NCCL; fallbacks, correctness under different collective algorithms, and elastic/fault-tolerant training are not addressed.
- Rounding modes and bias: rounding policy in FP8 conversions (e.g., RNE vs stochastic) and its influence on bias accumulation are not evaluated; potential benefits of stochastic rounding remain open.
- Error compensation: absence of error-feedback/residual accumulation mechanisms leaves potential accuracy gains unexplored; when error compensation helps/hurts with ASH+Dual-Scale is unclear.
- Distribution shifts and outliers: strategies for adapting τ and scales under abrupt activation distribution changes or outlier bursts (e.g., early training, curriculum shifts) are not provided.
- Overlap and scheduling contention: interactions between fused compression kernels and model compute (SM occupancy, register/shared memory pressure, stream priorities) are not characterized; potential regressions under high concurrency remain possible.
- Memory footprint and occupancy: shared memory/register usage vs occupancy trade-offs under varying B and model shapes lack profiling and guidelines; peak memory during fused ops is not reported.
- Interoperability with common training features: effects with activation checkpointing, activation offloading, ZeRO stage 2/3 optimizer sharding, and mixed-precision compute (BF16/FP16) are not measured.
- Baseline coverage: lack of TP-targeted compression baselines (e.g., low-rank/PowerSGD on TP tensors, sparsity/pruning of intermediate tensors) limits comparative understanding of TACO’s advantages and trade-offs.
- Reproducibility details: complete experimental protocols (datasets, token budgets, learning rate schedules, seeds), ablations per layer/operation, and public code availability are not specified, hindering independent validation.
- Assumption of Gaussianized post-transform distributions: whether ASH consistently induces near-Gaussian behavior across architectures, depths, and training phases, and when it fails (e.g., multimodal or heavy-tailed blocks), is untested.
- Alternative or learned transforms: potential gains from learned orthogonal transforms, per-layer transforms, or adaptive block sizing versus fixed FWHT are not explored.
- Failure handling and resilience: robustness to transient communication errors, metadata corruption, and recovery mechanisms (e.g., recomputation, checksums) in the compressed path is not discussed.
Practical Applications
Immediate Applications
The following applications can be put into practice now, leveraging the paper’s FP8-based TP communication compression (TACO), fused kernels, and integration with 3D parallel training.
- Sector: Software/AI infrastructure (LLM training at labs, cloud ML platforms, on-prem HPC)
- Use case: Drop-in acceleration of tensor-parallel (TP) LLM pretraining and finetuning
- What: Integrate TACO to compress TP intermediate tensors using FP8 with Adaptive Scale–Hadamard (ASH) and Dual-Scale (DS) quantization to reduce intra-node AllReduce/AllGather/ReduceScatter costs.
- Benefit: Up to ~1.87× TP throughput and ~1.53× across full 3D parallelism (TP+DP+PP) with near-lossless accuracy; higher cluster utilization and lower time-to-train.
- Tools/products/workflows:
- PyTorch/Megatron-LM plugin that intercepts TP collectives; a ProcessGroup/NCCL hook using COCCL two-shot AllReduce; fused CUDA kernels for compress/decompress.
- Configuration workflow: choose block size (e.g., shared-memory friendly sizes), enable COCCL two-shot AllReduce, set overlap chunk around 64 MB (as indicated), turn on ASH+DS FP8 TP path; integrate with SDP4Bit (DP) and TahQuant (PP) for 3D compression.
- Assumptions/dependencies:
- FP8 hardware support (e.g., NVIDIA H100/AMD MI300 or newer), NCCL/COCCL availability, ability to modify comm path in training stack, model uses TP. Validated on GPT models and Qwen2.5-7B; broader model families may require tuning.
- Sector: Cloud providers and MLOps/FinOps
- Use case: Cost and energy reduction SKUs for “compression-enabled training”
- What: Offer managed training configurations that enable FP8 TP compression with TACO for customers running Megatron/DeepSpeed/Megatron-Core stacks.
- Benefit: Lower GPU-hour costs and energy per training run; potential to offer “green tier” training options.
- Tools/products/workflows: Pre-baked AMIs/containers with COCCL + fused kernels; cost calculators showing throughput gains; policy guardrails to revert to higher precision on instability signals.
- Assumptions/dependencies: Customers must run TP with supported accelerators; monitoring hooks for divergence and validation loss; autoscaling/rollback logic.
- Sector: Enterprise/vertical AI (healthcare, finance, e-commerce, robotics)
- Use case: Faster foundation model pretraining and domain adaptation
- What: Apply TACO to speed up internal TP-heavy training (domain-specialized LLMs, code models, compliance models).
- Benefit: Accelerated R&D cycles; lower training bills; more frequent model refreshes for regulatory or market updates.
- Tools/products/workflows: Fine-tuning pipelines that toggle TACO; CI/CD training pipelines with regression checks on validation loss and downstream metrics.
- Assumptions/dependencies: Validation of accuracy in the target domain; careful monitoring for error accumulation; align with internal compliance requirements.
- Sector: Academic research and teaching
- Use case: Training larger models with limited budgets
- What: Use TACO to push TP-based training to larger widths/depths on fixed GPU budgets, and to study quantization/communication trade-offs.
- Benefit: More experiments per grant; reproducible studies of FP8 TP compression under different parallel splits.
- Tools/products/workflows: Open-source PyTorch extension with reference configs and ablation scripts; curriculum modules on comm-aware training and FP8 numerics.
- Assumptions/dependencies: Access to FP8-capable GPUs; availability of COCCL or alternative compressible-collective backends.
- Sector: Systems software (communication libraries, compilers, runtimes)
- Use case: Library-level support for compressible collectives
- What: Integrate TACO’s metadata layout (zero-copy αk, sk co-located with FP8 payload), fused kernels, and two-shot AllReduce into COCCL/NCCL plugins or MLIR/Triton codegen passes.
- Benefit: Turnkey performance improvements for frameworks; easier adoption by end users; consistent overlap schedules and chunk sizing.
- Tools/products/workflows: NCCL plugin module; Triton kernels for FWHT + quantization; CI benchmarks integrated with framework nightly builds.
- Assumptions/dependencies: ABI compatibility with NCCL; careful tuning for GPU SM/shared-memory constraints and warp-level reductions.
- Sector: Sustainability and ESG reporting (policy-adjacent operations)
- Use case: Immediate reduction of training energy and carbon per run
- What: Adopt TACO-enhanced training profiles for ongoing projects; report energy-per-token reductions as part of ESG disclosures.
- Benefit: Measurable decreases in power draw/time-to-train; aligns with internal sustainability goals and external reporting.
- Tools/products/workflows: Power/energy dashboards (e.g., DCGM/Prometheus + training profiler) showing compression-enabled efficiency.
- Assumptions/dependencies: Accurate power metering; controlled A/B comparisons; consistent dataset and schedule across runs.
- Sector: Model systems and architecture research
- Use case: Extending TP to multi-dimensional and hybrid parallel splits
- What: Apply TACO to TP + SP (sequence parallel) and expert-parallel (MoE) paths where intermediate tensors are frequently synchronized.
- Benefit: Better scaling within-node and reduced sensitivity to interconnect bottlenecks.
- Tools/products/workflows: MoE training stacks with compressed ReduceScatter/AllGather; autotuners that co-optimize expert routing with compression granularity.
- Assumptions/dependencies: Stability on MoE distributions; profiling-based tuning of block sizes and overlap windows.
Long-Term Applications
These opportunities require further research, ecosystem maturation, hardware features, or broader validation at scale.
- Sector: Hardware and interconnect co-design (semiconductors, data center networking)
- Use case: In-network or NIC-offloaded compressible collectives
- What: Offload ASH+DS FP8 compression/decompression and FWHT to SmartNICs/DPUs or switches; add hardware support for block-wise scaling and metadata handling.
- Benefit: Reduces GPU-side overhead further; enables compression across nodes without added latency.
- Tools/products/workflows: NIC firmware modules for two-shot AllReduce with compression; FPGA/ASIC FWHT blocks; CUDA-aware RDMA with metadata lanes.
- Assumptions/dependencies: Standards for FP8 formats and metadata; switch/NIC programmability; co-design with NCCL equivalents.
- Sector: ML compilers and graph optimization
- Use case: End-to-end compiler integration and auto-scheduling of compression
- What: Compile-time placement of compress/decompress ops, automatic block-size selection, and overlap chunk scheduling based on graph shape and hardware counters.
- Benefit: Removes manual tuning; portable performance across model families and clusters.
- Tools/products/workflows: TVM/MLIR/Triton passes that treat TP collectives as optimization sites; feedback-driven autotuners using runtime telemetry.
- Assumptions/dependencies: Stable APIs for collective interception; robust performance models for bandwidth/latency and kernel fusion limits.
- Sector: Generalization to broader model classes
- Use case: Multimodal, graph, and RL training with high-frequency intermediate synchronization
- What: Extend TACO to ViTs, diffusion, speech, GNNs, and RL policy/value networks with TP-like communication patterns.
- Benefit: Cross-domain training acceleration; common compression substrate in heterogeneous workloads.
- Tools/products/workflows: Domain-aware ASH targets (τ) and block partitioning strategies; per-layer/pipeline policy for when to compress.
- Assumptions/dependencies: Distributional differences vs. LLM activations; convergence studies; sensitivity to optimizer/regularization.
- Sector: Robust, adaptive compression controllers
- Use case: Closed-loop control of precision and scales to prevent divergence
- What: Online adaptation of αk, sk, FP8 variant (E4M3 vs E5M2), and fallback thresholds based on validation loss spikes and gradient statistics.
- Benefit: Maintains stability across long pretraining runs; minimizes human intervention.
- Tools/products/workflows: Controllers embedded in training loops; anomaly detectors that trigger temporary de-compression (e.g., FP16) for sensitive layers/steps.
- Assumptions/dependencies: Reliable online metrics; low-latency control paths; standardized escalation policies.
- Sector: Standards and governance (policy, industry consortia)
- Use case: Best-practice guidelines and benchmarks for communication-compressed training
- What: Define test suites for accuracy stability under FP8 TP compression, publish energy/throughput metrics, and codify reporting norms.
- Benefit: Comparable, transparent efficiency claims; safer adoption by public and regulated sectors.
- Tools/products/workflows: MLPerf-style tracks for comm-compressed 3D parallel training; procurement checklists requiring energy-per-token and stability audits.
- Assumptions/dependencies: Community consensus on metrics; participation from hardware vendors and framework maintainers.
- Sector: Training economics and capacity planning
- Use case: Cluster-level schedulers aware of compression-induced speedups
- What: Job admission/control that predicts runtime and energy with TACO; scheduled chunk overlaps aligned with network topology.
- Benefit: Higher cluster throughput and fairness; improved ROI on accelerators.
- Tools/products/workflows: Scheduler plugins that simulate two-shot AllReduce with compression; cost models integrated with FinOps dashboards.
- Assumptions/dependencies: Accurate performance models; telemetry from production runs; multi-tenant isolation considerations.
- Sector: Security and reliability
- Use case: Fault-tolerant compressed collectives and drift detection
- What: Integrate compression with error-detecting codes and checksums; detect scale/metadata corruption; resilient recovery without retraining epochs.
- Benefit: Production robustness; reduces costly restarts due to silent accuracy drift.
- Tools/products/workflows: Redundant metadata lanes; periodic high-precision sanity checks; replay buffers for critical steps.
- Assumptions/dependencies: Overhead budget for verification; compatible with comm libraries and checkpointing.
- Sector: Consumer and societal impact (indirect)
- Use case: Faster model iteration cycles for downstream apps (search, assistants, education, healthcare triage)
- What: Organizations can refresh models more often or broaden access via lower-cost training.
- Benefit: Improved product quality and availability; potential democratization of advanced models.
- Tools/products/workflows: Shorter release cadences; expanded A/B testing coverage with compressed training runs.
- Assumptions/dependencies: Benefits depend on how providers pass savings to users; safeguards for accuracy and safety maintained under faster iteration.
Notes on feasibility across applications:
- TACO’s stability and accuracy are demonstrated on GPT-350M and Qwen2.5-7B with modern GPUs and high-bandwidth intra-node interconnects; extrapolation to trillion-parameter or highly heterogeneous, multi-node regimes requires validation.
- Effectiveness depends on model distributions matching assumptions (dense near-zero intermediate tensors); re-tuning τ, block sizes, and FP8 variant may be needed across architectures.
- Compression frequency and algorithm (two-shot AllReduce) materially affect both performance and numerical robustness; deviations from these protocols can erode gains.
- Monitoring and fallback mechanisms are recommended to mitigate rare but impactful training instabilities.
Glossary
- 3D parallelism: A training strategy that combines data, tensor, and pipeline parallelism to scale large models across many devices; "3D-parallel training systems"
- Adaptive Scale--Hadamard Transform: A block-wise scaling plus orthogonal Hadamard rotation used to reshape tensor distributions for more faithful low-precision quantization; "Adaptive Scale--Hadamard Transform to dynamically reshape the distribution of TP intermediate tensors"
- AllGather: A collective communication primitive that gathers data from all ranks and distributes the concatenated result to all; "decomposed into AllGather and Reduce-Scatter operations"
- AllReduce: A collective operation that reduces values (e.g., sums) across all ranks and broadcasts the result back; "including the AllReduce operations required during forward and backward passes."
- AlltoAll: A collective where each rank sends distinct data to every other rank and receives distinct data from each; "one compressed AlltoAll operation"
- ASH transform: The abbreviation for the Adaptive Scale–Hadamard Transform used in TACO’s pipeline; "followed by the ASH transform"
- coalesced global memory accesses: A GPU memory access pattern where threads access contiguous addresses to maximize bandwidth utilization; "enables fully coalesced global memory accesses during reconstruction"
- collective communication: Coordinated communication operations among multiple processes/GPUs (e.g., AllReduce, AllGather); "synchronized via collective communication"
- COCCL: A compressible collective communication library used to integrate compression with NCCL collectives; "COCCL, a compressible collective communication library"
- CTA (Cooperative Thread Array): NVIDIA’s term for a CUDA thread block that executes on a single SM; "parallel CTAs."
- Data parallelism (DP): A parallel training strategy that replicates the model across workers and synchronizes gradients; "data parallelism (DP)"
- Dual-Scale Quantization: A scheme that uses two scales per block—one for adaptive rescaling and one for post-rotation quantization—to ensure stability and dynamic range; "implements Dual-Scale Quantization to keep all communicated values within the FP8 representable range"
- E4M3: An FP8 format with 4 exponent bits and 3 mantissa bits, trading dynamic range and precision; "Common FP8 variants include E4M3 and E5M2."
- E5M2: An FP8 format with 5 exponent bits and 2 mantissa bits, offering a different precision–range tradeoff; "Common FP8 variants include E4M3 and E5M2."
- EDGC: Entropy-Driven Dynamic Gradient Compression, a method to compress gradients based on entropy to reduce latency; "EDGC uses entropy-driven dynamic gradient compression to reduce latency in GPT training"
- Fast Walsh--Hadamard Transform (FWHT): An O(n log n) algorithm to apply the Walsh–Hadamard transform efficiently; "implemented using the Fast Walsh--Hadamard Transform (FWHT)"
- FlatQuant: An inference-oriented quantization method cited as related work on low-bit compression; "SmoothQuant, FlatQuant, QuaRot, and GPTQ"
- FP8: An 8-bit floating-point format (e.g., E4M3/E5M2) used to reduce communication/memory with exponent scaling; "a robust FP8-based framework for compressing TP intermediate tensors"
- GPTQ: A post-training quantization method for transformers, listed among inference-oriented techniques; "SmoothQuant, FlatQuant, QuaRot, and GPTQ"
- Hadamard transform: An orthogonal transformation using Hadamard matrices that redistributes signal energy; "the standard Hadamard transform"
- INT8: 8-bit integer quantization with uniform steps and optional zero-point, often used for inference; "INT8 quantization typically adopts a fixed scale with zero-point"
- kernel fusion: Combining multiple GPU operations into a single kernel to reduce memory traffic and launch overhead; "with kernel fusion and communication–computation co-optimization"
- kernel launch overhead: The performance cost incurred when starting GPU kernels, especially numerous small launches; "reduce memory traffic and kernel launch overhead"
- LLM-QAT: A quantization-aware training method tailored for LLMs; "Quantization-aware training methods, such as LLM-QAT"
- NCCL: NVIDIA Collective Communications Library for high-performance multi-GPU collectives; "built on NCCL."
- Pipeline parallelism (PP): A parallelism technique that partitions model layers across devices and pipelines microbatches; "pipeline parallelism (PP)"
- quantization-aware training: Training that models quantization effects during forward/backward to improve robustness; "Quantization-aware training methods"
- QuaRot: A quantization approach (e.g., rotation-based) listed among inference-oriented methods; "SmoothQuant, FlatQuant, QuaRot, and GPTQ"
- Reduce-Scatter: A collective that reduces arrays across all ranks and scatters disjoint reduced chunks to each rank; "decomposed into AllGather and Reduce-Scatter operations"
- ReduceScatter phase: The stage of an AllReduce decomposition where reduction and scattering are performed; "The ReduceScatter phase consists of one compressed AlltoAll operation"
- root mean square (RMS): A measure of signal magnitude used here for block-wise scaling; "root mean square (RMS) amplitude"
- SDP4bit: A 4-bit gradient communication method optimized for sharded data-parallel training; "SDP4bit achieves near-4-bit communication in sharded DP training"
- Sequence parallelism (SP): A technique that parallelizes across the sequence dimension, affecting TP communication patterns; "sequence parallelism (SP)"
- SmoothQuant: A technique for smoothing activation/weight distributions to enable low-bit inference; "SmoothQuant, FlatQuant, QuaRot, and GPTQ"
- subnormal range: The range of floating-point numbers smaller than the smallest normal, with reduced precision; "fall into the subnormal range"
- TahQuant: A pipeline-parallel activation quantization method used as a baseline; "TahQuant introduces fine-grained activation quantization along the PP communication path"
- TAGC: Transformer-Aware Gradient Compression, a hierarchical compression scheme for transformers; "TAGC introduces a Transformer-aware hierarchical compression scheme"
- Tango: A system that coordinates quantization with scheduling to reduce training overhead; "e.g., Tango, which reduces computation and communication overhead through quantization-aware scheduling"
- Tensor parallelism (TP): Partitioning tensor dimensions (e.g., hidden rows/columns) across devices with frequent synchronization; "tensor parallelism (TP)"
- two-shot AllReduce: An AllReduce algorithm decomposed into ReduceScatter and AllGather steps to optimize performance; "using COCCL’s two-shot AllReduce as the communication algorithm"
- Unit in the last place (ULP): The spacing between adjacent representable floating-point numbers at a given magnitude; "unit in the last place (ULP)"
- Walsh--Hadamard transform: An orthogonal transform with ±1 entries used for energy redistribution; "an orthogonal Walsh--Hadamard transform:"
- warp-level reductions: Intra-warp aggregation using shuffle or similar primitives for fast parallel reductions; "reuses warp-level reductions"
- zero-point: An offset used in uniform integer quantization to map zero exactly; "fixed scale with zero-point"
Collections
Sign up for free to add this paper to one or more collections.