# Data-Parallel Training

ggmlR provides `dp_train()` for data-parallel training across multiple GPUs (or CPU cores). Each replica processes a unique sample per step; gradients are averaged and applied once.

    library(ggmlR)

## 1. Concept

`dp_train()` takes a model factory (`make_model`) instead of a model instance. It creates `n_gpu` identical replicas, synchronises their initial weights, and runs a gradient-accumulation loop:

    for each iteration:
      each replica  →  forward(sample_i)  →  loss  →  backward
      average gradients across replicas
      optimizer step on replica 0
      broadcast updated weights to all replicas

The effective batch size equals `n_gpu` (one sample per replica per step).

## 2. Minimal example

    data(iris)
    set.seed(42)

    x_cm <- t(scale(as.matrix(iris[, 1:4])))    # [4, 150]
    y_oh <- t(model.matrix(~ Species - 1, iris)) # [3, 150]

    # Dataset as list of (x, y) pairs — one sample each
    dp_data <- lapply(seq_len(ncol(x_cm)), function(i)
      list(x = x_cm[, i, drop = FALSE],
           y = y_oh[, i, drop = FALSE]))

    # Model factory — called once per replica
    make_model <- function() {
      ag_sequential(
        ag_linear(4L, 32L, activation = "relu"),
        ag_linear(32L, 3L)
      )
    }

    result <- dp_train(
      make_model = make_model,
      data       = dp_data,
      loss_fn    = function(out, tgt) ag_softmax_cross_entropy_loss(out, tgt),
      forward_fn = function(model, s)  model$forward(ag_tensor(s$x)),
      target_fn  = function(s)         s$y,
      n_gpu      = 1L,         # set to ggml_vulkan_device_count() for multi-GPU
      n_iter     = 2000L,
      lr         = 1e-3,
      verbose    = TRUE
    )

    cat("Final loss:", result$loss, "\n")
    model <- result$model

## 3. Multi-GPU

    n_gpu <- max(1L, ggml_vulkan_device_count())
    cat(sprintf("Training on %d GPU(s)\n", n_gpu))

    result_mg <- dp_train(
      make_model = make_model,
      data       = dp_data,
      loss_fn    = function(out, tgt) ag_softmax_cross_entropy_loss(out, tgt),
      forward_fn = function(model, s)  model$forward(ag_tensor(s$x)),
      target_fn  = function(s)         s$y,
      n_gpu      = n_gpu,
      n_iter     = 2000L,
      lr         = 1e-3,
      max_norm   = 5.0,      # gradient clipping
      verbose    = FALSE
    )

With `n_gpu = 2` the effective batch is 2 and training is ~2x faster (ignoring communication overhead).

## 4. Gradient clipping

Pass `max_norm` to clip the global gradient norm before each optimizer step:

    result <- dp_train(
      make_model = make_model,
      data       = dp_data,
      loss_fn    = function(out, tgt) ag_softmax_cross_entropy_loss(out, tgt),
      forward_fn = function(model, s)  model$forward(ag_tensor(s$x)),
      target_fn  = function(s)         s$y,
      n_gpu      = 1L,
      n_iter     = 2000L,
      lr         = 1e-3,
      max_norm   = 1.0       # clip to unit norm
    )

## 5. `ag_dataloader` — batched training loop

For standard single-process batched training `ag_dataloader` is simpler than `dp_train`:

    x_tr <- x_cm[, 1:120];  y_tr <- y_oh[, 1:120]

    dl <- ag_dataloader(x_tr, y_tr, batch_size = 32L, shuffle = TRUE)

    model2  <- make_model()
    params2 <- model2$parameters()
    opt2    <- optimizer_adam(params2, lr = 1e-3)

    ag_train(model2)
    for (ep in seq_len(100L)) {
      for (batch in dl$epoch()) {
        with_grad_tape({
          loss <- ag_softmax_cross_entropy_loss(
            model2$forward(batch$x), batch$y$data)
        })
        grads <- backward(loss)
        opt2$step(grads);  opt2$zero_grad()
      }
    }

## 6. Full example

A detailed example with synthetic regression, multiple replica counts, and correctness checks:

    # inst/examples/dp_train_demo.R

Multi-GPU scheduler usage (low level):

    # inst/examples/multi_gpu_example.R