Data Loading Optimization — Hyperdisk ML, Parallelstore, Volume Populator

Tại sao data loading là critical path

Training large models:

Timeline (Epoch 1, 1M steps):
  Step 1: Load batch (100ms) → GPU compute (900ms)
  Step 2: Load batch (100ms) → GPU compute (900ms)
  ...
  Step 1M: Load batch (100ms) → GPU compute (900ms)
  
Total: 1M × (100ms + 900ms) = 1,000M ms = 278 hours

But if slow storage (10GB Cloud Storage direct read):
  Step 1: Load batch (5000ms!) → GPU compute (900ms)
  Step 2: Load batch (5000ms) → GPU compute (900ms)
  ...
  Total: 1M × (5000ms + 900ms) = 5,900M ms = 1,638 hours
  
Slowdown: 5.9x due to data loading alone!

Implication: Without optimized data pipeline, GPU compute is 90% idle waiting for data.

---

Data loading bottlenecks

Source bottleneck: Cloud Storage (GCS)

Direct read from GCS:

10,000 small files (10KB-100KB each)
Training loop: for file in files: model(load(file))

Per-file latency: 100ms (metadata lookup + download)
Total: 10,000 × 100ms = 1,000 seconds per epoch
GPU: 99% idle waiting for data

Network bottleneck: Serial download

Model wants 64 files per batch
Sequential download:
  File 1: 100ms
  File 2: 100ms
  ...
  File 64: 100ms
  Total: 6.4 seconds per batch

Parallel download (8 workers):
  Batch 1-8: parallel (100ms)
  Total: 1.2 seconds per batch
  
Speedup: 5.3x from parallelization alone

Storage latency: Random access

Training framework: TensorFlow dataset.shuffle() (random access)
Each sample: random seek in storage system

HDD (spinning disk): 5-50ms per seek → terrible
SSD (local): <1ms per seek → good
Network SSD (NVMe): 1-10ms (network latency) → mediocre

---

Storage solutions: Hyperdisk ML

Characteristics

Hyperdisk ML is block storage (not filesystem):

Hyperdisk ML (block storage):
  Throughput: up to 64 GB/s (aggregate across VMs)
  Latency: ~1-5ms (network SSD)
  Cost: ~$0.10/GB-month
  Access: RWX (read-write-many, concurrent access)

Standard Persistent Disk:
  Throughput: 10-20 GB/s
  Latency: 5-10ms
  Cost: ~$0.04/GB-month
  Access: RWO (read-write-one, exclusive)

Parallelstore:
  Throughput: >500 GB/s (distributed filesystem)
  Latency: sub-millisecond
  Cost: ~$0.30/GB-month
  Access: POSIX filesystem (RWX)

Model weight loading with Hyperdisk ML

Scenario: Load 70B LLM weights (140GB model, fp32)

With GCS direct:

Sequential read: 140GB / 1 GB/s = 140 seconds per load
5 model replicas: 5 × 140s = 700 seconds total = 11 minutes

With Hyperdisk ML (RWX, 64 GB/s aggregate):
Load once on shared volume, 5 Pods mount:
  1 Pod downloads from GCS → Hyperdisk (140GB)
  Time: 140GB / 2 GB/s (GCS read) = 70 seconds
  
  5 Pods read from Hyperdisk simultaneously:
  Time: 140GB / 64 GB/s = 2.2 seconds per Pod
  Total: 70s (GCS download) + 2.2s (local read) = 72 seconds
  
  vs 700 seconds: 10x faster!

GKE setup with Hyperdisk ML

# 1. Create Hyperdisk ML volume
gcloud compute disks create model-weights \
  --size=150GB \
  --type=hyperdisk-balanced \
  --replica-zones=us-central1-a,us-central1-b

# 2. Create PVC (read-only-many)
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: model-weights-pvc
spec:
  accessModes:
  - ReadOnlyMany  # Multiple Pods can read concurrently
  storageClassName: hyperdisk-storage
  resources:
    requests:
      storage: 150Gi

# 3. Pods mount volume
---
apiVersion: v1
kind: Pod
metadata:
  name: inference-replica-1
spec:
  containers:
  - name: inference
    volumeMounts:
    - name: model-weights
      mountPath: /model
  volumes:
  - name: model-weights
    persistentVolumeClaim:
      claimName: model-weights-pvc

---

Parallelstore: Distributed filesystem for HPC/AI

Cơ chế

Parallelstore is managed distributed filesystem:

Single namespace (e.g., /data/) distributed across:
  - Multiple servers
  - Multiple disks per server
  - Striped data (RAID-like, fault-tolerant)

Characteristics:
  Throughput: linear scaling with clients
              1 client: 50 GB/s
              10 clients: 500 GB/s
  Latency: sub-millisecond (optimized for AI/HPC)
  Metadata: fast (optimized for small file access)

Parallelstore for training data

Scenario: ImageNet (1.2M images, 150GB total)

With GCS (10k files):

Training loop: dataset.shuffle() + load images
Sequential: 50ms per image × 1M/epoch = 50M seconds = 14 hours wasted on I/O

With Parallelstore:
Sequential: 0.5ms per image × 1M/epoch = 0.5M seconds = 139 seconds
GPU: 99.6% utilized instead of 5%

GKE setup with Parallelstore

# 1. Create Parallelstore instance
gcloud parallelstore instances create training-data \
  --location=us-central1-a \
  --capacity=200GB \
  --tier=high-performance

# 2. Mount in GKE via CSI driver
gcloud container clusters update my-cluster \
  --enable-parallelstore=true

# 3. Pod mounts Parallelstore volume
---
apiVersion: v1
kind: Pod
metadata:
  name: training-job
spec:
  containers:
  - name: trainer
    volumeMounts:
    - name: training-data
      mountPath: /datasets
  volumes:
  - name: training-data
    persistentVolumeClaim:
      claimName: training-data-pvc

Comparison: Hyperdisk ML vs Parallelstore

Aspect	Hyperdisk ML	Parallelstore
Throughput	64 GB/s aggregate	>500 GB/s
Latency	1-5ms	<1ms
Model weights	✓ Excellent	✓ Good
Training data (many small files)	✗ Mediocre	✓ Excellent
Cost per GB	Lower	Higher
POSIX filesystem	✗ Block device	✓ Full filesystem
Scaling	Fixed capacity	Grows on-demand

Decision:

• Model weights (few large blobs) → Hyperdisk ML
• Training data (many small files) → Parallelstore

---

GKE Volume Populator: Automated data prep

Problem: Cold start

Training job starts, needs 200GB training data
Option 1: Download from GCS every epoch
  Epoch 1: download 200GB + train → 3 hours
  Epoch 2: re-download 200GB + train → 3 hours
  ...
  Total: 10 epochs × 3 hours = 30 hours
  
Option 2: Pre-download data once
  Setup: download 200GB → Hyperdisk ML (takes 1 hour)
  Epoch 1: read from Hyperdisk (fast) → 1 hour
  Epoch 2: read from Hyperdisk (fast) → 1 hour
  ...
  Total: 1 hour (setup) + 10 hours (training) = 11 hours
  
Savings: 19 hours (63%) by pre-staging data

Volume Populator (automatic prep)

GKE Volume Populator automatically:

1. User specifies: "populate volume from GCS path"
2. GKE creates Hyperdisk ML/Parallelstore volume
3. GKE launches background job: download data from GCS
4. When download complete: training Pod can mount + use
5. On teardown: volume auto-deleted

Setup

# Volume populator via DataSourceRef
---
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: training-data
spec:
  dataSourceRef:
    apiGroup: kubernetes.io
    kind: PersistentVolumeClaim
    name: gcs-source-data  # Reference to GCS data source
  
  accessModes:
  - ReadWriteMany
  storageClassName: parallelstore-fast
  resources:
    requests:
      storage: 200Gi

---
# GCS data source definition
apiVersion: v1
kind: ConfigMap
metadata:
  name: gcs-source-data
data:
  source-uri: "gs://my-bucket/training-data/"
  # Volume populator reads this and triggers download

GKE workflow:

t=0: PVC created
t=0-1min: GKE launches populator job (copy gs:// → Parallelstore)
t=1-30min: Data download (200GB over network)
t=30min: Volume ready
t=30min+: Training Pod mounts volume → starts immediately

---

Data pipeline optimization strategies

Strategy 1: Prefetching

# TensorFlow dataset optimization
dataset = tf.data.Dataset.from_tensor_slices(file_paths)
dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
# Load next batch while GPU computes current batch

Strategy 2: Caching

# Cache first epoch to memory/disk
dataset = dataset.cache('/tmp/dataset_cache')
# Subsequent epochs read from cache (fast)

Strategy 3: Pipelining

# Parallelize data loading
dataset = dataset.interleave(
    tf.data.TFRecordDataset,
    cycle_length=8,  # 8 files in parallel
    block_length=16
)

Strategy 4: Format optimization

CSV (unstructured):
  Parsing overhead: high
  Compression: poor
  Read speed: slow

TFRecord (binary):
  Parsing overhead: low (pre-serialized)
  Compression: good
  Read speed: fast (up to 10x faster than CSV)

---

Cost vs performance trade-off

Training 70B LLM, 10 epochs (1 week wallclock):

Option A: GCS direct
  Compute cost: 10 GPUs × 7 days × $15/hr = $25,200
  Storage cost: 0 (GCS is "free")
  Total: $25,200
  Wallclock: 7 days (slow I/O)

Option B: Hyperdisk ML pre-staging
  Compute cost: 10 GPUs × 1 day (1 extra day for setup)
  Hyperdisk cost: 150GB × $0.10/month = $1.50
  Storage cost: <<$1
  Total: $25,200 + $1.50 = $25,201.50
  Wallclock: 6 days (fast I/O)
  
Savings: $0 saved, but 1 day faster (compute time same, but earlier completion)

Option C: Parallelstore (expensive but fast)
  Compute cost: 10 GPUs × 0.5 day (2x faster I/O) × $15/hr = $1,800
  Parallelstore cost: 200GB × $0.30/month = $60
  Total: $1,860 (compute only, not wallclock)
  Wallclock: 3 days (very fast I/O)
  
Break-even: if job is deadline-critical (reduce wallclock), Parallelstore worth it
            if job is cost-critical, use Hyperdisk ML (cheap overhead)

---

Mental model: Data pipeline design

Step 1: Profile bottleneck
  → GPU waiting for data? (low GPU utilization, high I/O wait)
  → Or GPU computing? (high GPU utilization)

Step 2: Match storage tier
  Model weights (sequential, large) → Hyperdisk ML
  Training data (random, small files) → Parallelstore
  Temporary cache → local node SSD

Step 3: Configure pipeline
  Prefetch + cache + interleave (framework-level)
  Volume Populator (GKE-level)

Step 4: Monitor
  GPU utilization should be >90%
  If <80% → data pipeline still bottleneck

---

References

• GKE Hyperdisk ML
• GKE Parallelstore Optimization
• GKE Volume Populator
• GKE Parallelstore CSI Driver
• TensorFlow Dataset Best Practices