---
title: TrainingCXL Architecture
url: https://www.emergentmind.com/topics/trainingcxl-architecture
type: topic
---

# TrainingCXL Architecture

TrainingCXL is an architecture designed to facilitate efficient, fault-tolerant training of large-scale recommendation models by leveraging persistent memory (PMEM) disaggregation over Compute Express Link (CXL). It achieves high throughput and energy efficiency by tightly integrating PMEM with GPU memory hierarchies, offloading embedding computation and checkpointing logic to on-card hardware near the CXL controller, and introducing advanced checkpointing techniques that relax parameter-update dependencies across batches. Key reported outcomes include a 5.2× performance improvement and 76% energy savings versus conventional PMEM-based recommendation system training [2301.07492].

## 1. System Architecture and Hardware Composition

TrainingCXL is built around a Type-2 CXL architecture where both CXL-enabled GPUs (CXL-GPU) and disaggregated PMEM-based memory expanders (CXL-MEM) attach to the same CXL root complex via CXL switches. Both are Type-2 devices and interact over CXL.mem and CXL.cache protocols, sharing a single Host Physical Address (HPA) space. The Device Coherency Engine (DCOH) enables direct cache-coherent access between GPU and memory expander, so the GPU can issue load/store requests directly to PMEM-resident data, eliminating the need for CPU-mediated communication or software copying.

This integration makes embeddings and model parameters stored in PMEM appear as extensions of the GPU’s memory hierarchy, allowing the GPU to operate on large recommendation datasets and models with minimal data movement overhead.

## 2. On-Card Computing and Checkpointing Logic

The CXL-MEM device hosts two dedicated hardware modules behind its CXL endpoint controller:

- **Computing Logic:** Includes scratchpad memory for buffering embedding vectors, arrays of adders and multipliers to implement "lookup & update" primitives for sparse embedding tables, and a microsequencer configured via MMIO-programmed registers that manage parameters such as embedding length and learning rates.
  
- **Checkpointing Logic:** Features a lightweight DMA engine that can generate CXL.mem read/write requests and CXL.cache fetches (for fetching MLP parameters from CXL-GPU), two 32-bit counters for embedding and MLP logs, a persistent-flag per log entry to track log durability, and a small state-machine responsible for log management including undo logging, flagging, and garbage collection of obsolete checkpoints.

The data/control flow involves host-initiated MMIO writes to CXL-MEM specifying affected embedding indices and parameter addresses, background embedding lookup, concurrent undo logging to PMEM, MLP parameter logging post-GPU update, and automated garbage collection once durability is confirmed via persistent flags.

## 3. Persistency Mechanisms and Fault Tolerance

Persistency is enforced through a region-divided PMEM layout, separating the "data region" (live embedding tables) and "log region" (undo logs for embeddings and MLP parameters). 

The batch-aware undo logging protocol operates as follows: prior to updates, the host communicates which embedding indices will change; CXL-MEM then copies pre-update vectors into the log region, marking durability by setting persistent flags. Only after this flag is set are writes permitted to the data region. MLP parameters are fetched from GPU post-compute and similarly logged.

Checkpointing occurs off the training critical path—logging of embedding and MLP data overlaps with GPU compute phases (Bottom-MLP, feature interaction, Top-MLP), masking log latency and ensuring failure atomicity even under PMEM failure or host disruptions.

## 4. Relaxed Checkpointing Across Batches

TrainingCXL applies two main relaxations to erase remaining stalls from RAW hazards and logging overhead:

### a) Relaxed Embedding Lookup

Empirical observation shows that 80% of embeddings accessed in batch N+1 were updated in batch N, resulting in frequent PMEM RAW dependencies with traditional update-order semantics. TrainingCXL uses "commute-reorder": lookups for batch N+1 are performed using the snapshot of the embedding table at batch N, and only subsequently are updates for batch N applied. The accumulation for the new batch is finalized after updates, effectively decoupling dependency chains.

```python
for batch = 1 to B:
    lookup_buf ← PMEM.read(table_addr_old, indices[batch])
    # overlap: CXL-MEM does undo-log if not yet done
    table_addr_new[indices[batch]] += g  # gradient update
    finalize_buf(lookup_buf)
```

### b) Relaxed MLP Checkpointing

Instead of checkpointing MLP weights every batch, logs are batched across k batches (e.g., \(k=100\)), based on the finding that such periodicity maintains 0.01% accuracy-loss requirements. During these k batches, CXL-MEM suspends new logs when the GPU Top-MLP phase starts, fully overlapping log creation with existing compute, driving the net checkpoint overhead to near zero.

## 5. Analytical Performance and Energy Model

The batch latency \(T_{\mathrm{batch}}\) is defined as:

\[
T_{\mathrm{batch}} = \max\left\{
  T_{\mathrm{BMLP}},
  T_{\mathrm{EmbLookup}},
  T_{\mathrm{EmbUpdate}} + T_{\mathrm{UndoLog}},
  T_{\mathrm{MLPLog}}
\right\} + T_{\mathrm{FI}} + T_{\mathrm{TMLP}}
\]

where each term represents the respective stage: Bottom-MLP (\(T_{\mathrm{BMLP}}\)), embedding lookup/update, undo log, MLP log, feature interaction (\(T_{\mathrm{FI}}\)), and Top-MLP (\(T_{\mathrm{TMLP}}\)). TrainingCXL’s overlap and dependency relaxation reduce this to:

\[
T_{\mathrm{batch}}^{\mathrm{CXL}} \approx \max\{T_{\mathrm{BMLP}},T_{\mathrm{EmbLookup}}^{\mathrm{relaxed}}\} + T_{\mathrm{FI}} + T_{\mathrm{TMLP}}
\]

since embedding update/logging and MLP logging are hidden behind compute.

TrainingCXL achieves a measured speedup:

\[
\mathrm{Speedup} = \frac{T_{\mathrm{batch}}^{\mathrm{PMEM}}}{T_{\mathrm{batch}}^{\mathrm{CXL}}} \approx 5.2\times
\]

Energy per batch:

\[
E_{\mathrm{batch}} = P_{\mathrm{GPU}}\cdot T_{\mathrm{GPU_{compute}}} + P_{\mathrm{PMEM}}\cdot T_{\mathrm{PMEM_{mem\,I/O}}} + P_{\mathrm{CXL}}\cdot T_{\mathrm{CXL\,logic_{oncard\,logic}}}
\]

Reported overlap strategies reduce \(T_{\mathrm{PMEM}}\) and \(T_{\mathrm{CXL\,logic}}\) on the critical path, yielding 76% energy savings over PMEM-only architectures [2301.07492].

## 6. Architectural Figures and Prototyping

The original presentation includes a range of schematics:

- CXL integration and topology (CXL.mem, CXL.cache among host, GPU, and memory expander)
- Block diagrams of TrainingCXL’s hardware, including CXL switches, PMEM modules, and the detailed floor-plan of FPGA-based prototypes
- Data/control flow diagrams illustrating the sequence of embedding and MLP parameter log creation, automated hardware-managed checkpointing, and commutative-relaxed lookup/update schemes across batches

Together, these illustrate the full TrainingCXL stack: system architecture, hardware logic, persistency workflow, and batch-straddling checkpointing protocol that collectively realize efficient, software-transparent, and failure-tolerant training for large-scale recommendation models with substantial performance and energy advantages [2301.07492].

Source: https://www.emergentmind.com/topics/trainingcxl-architecture