server /train: examples rendered through the chat template, one window each, loss only on the assistant turns (+ the CE backward honouring masked rows)

common/common.cpp

@@ -2136,6 +2136,27 @@ ggml_opt_dataset_t common_opt_dataset_init(struct llama_context * ctx, const std
     return result;
 }
 
+ggml_opt_dataset_t common_opt_dataset_init_masked(int64_t window, const std::vector<std::vector<llama_token>> & seqs,
+                                                  const std::vector<std::vector<uint8_t>> & loss, llama_token pad) {
+    GGML_ASSERT(seqs.size() == loss.size());
+    const int64_t ndata = (int64_t) seqs.size();
+    ggml_opt_dataset_t result = ggml_opt_dataset_init(GGML_TYPE_I32, GGML_TYPE_I32, window, window, ndata, /*ndata_shard =*/ 1);
+
+    llama_token * data   = (llama_token *) ggml_opt_dataset_data(result)->data;
+    llama_token * labels = (llama_token *) ggml_opt_dataset_labels(result)->data;
+
+    for (int64_t idata = 0; idata < ndata; ++idata) {
+        const auto & s = seqs[idata];
+        const auto & m = loss[idata];
+        GGML_ASSERT(s.size() == m.size() && (int64_t) s.size() <= window + 1);
+        for (int64_t j = 0; j < window; ++j) {
+            data  [idata*window + j] = j < (int64_t) s.size() ? s[j] : pad;
+            labels[idata*window + j] = (j + 1 < (int64_t) s.size() && m[j + 1]) ? s[j + 1] : -1;
+        }
+    }
+    return result;
+}
+
 ggml_opt_optimizer_params common_opt_lr_pars(void * userdata) {
     ggml_opt_optimizer_params result = ggml_opt_get_default_optimizer_params(nullptr);
     const lr_opt &            d      = *(lr_opt *) userdata;

common/common.h

@@ -1135,6 +1135,13 @@ inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() {
 
 ggml_opt_dataset_t common_opt_dataset_init(struct llama_context * ctx, const std::vector<llama_token> & tokens, int64_t stride);
 
+// One datapoint per sequence: seqs[i] (at most window + 1 tokens) is the input, left-aligned and
+// padded with `pad` to `window`; the label at position j is seqs[i][j + 1] where loss[i][j + 1]
+// is set, and -1 (no loss) everywhere else, padding included. Completion-only training marks
+// only the assistant's tokens.
+ggml_opt_dataset_t common_opt_dataset_init_masked(int64_t window, const std::vector<std::vector<llama_token>> & seqs,
+                                                  const std::vector<std::vector<uint8_t>> & loss, llama_token pad);
+
 // "adamw" or "sgd" (case insensitive)
 enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *);
 

ggml/src/ggml-cpu/ops.cpp

@@ -11884,7 +11884,14 @@ static void ggml_compute_forward_cross_entropy_loss_back_f32(
         assert(sum > 0.0);
         ggml_vec_scale_f32(nc, ds0, 1.0/sum);
 
-        // grad(src0f) = (softmax(src0f) - src1f) * grad(cross_entropy_loss(src0f, src1f)) / nr
+        // grad(src0f) = (softmax(src0f) * sum(src1f) - src1f) * grad(cross_entropy_loss(src0f, src1f)) / nr
+        // The derivative of -sum(y log softmax(x)) is softmax(x) * sum(y) - y. Dropping sum(y) is
+        // exact only when every label row sums to 1; a MASKED row (all zero: a position that must
+        // carry no loss, e.g. a prompt token in completion-only training) would still push
+        // softmax(x) back through the graph. With the sum the masked row's gradient is zero.
+        float lsum = 0.0f;
+        ggml_vec_sum_f32(nc, &lsum, s1);
+        ggml_vec_scale_f32(nc, ds0, lsum);
         ggml_vec_sub_f32(nc, ds0, ds0, s1);
         ggml_vec_scale_f32(nc, ds0, d_by_nr);
 

ggml/src/ggml-cuda/cross-entropy-loss.cu

@@ -82,7 +82,15 @@ static __global__ void cross_entropy_loss_back_f32(
         }
     }
     sum = warp_reduce_sum(sum);
-    const float sm_scale = 1.0f/sum;
+
+    // d/dx of -sum(y log softmax(x)) is softmax(x)*sum(y) - y: with sum(y) a masked (all-zero)
+    // label row gets a zero gradient; without it the row would still push softmax(x) back.
+    float lsum = 0.0f;
+    for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) {
+        lsum += labels[i];
+    }
+    lsum = warp_reduce_sum(lsum);
+    const float sm_scale = lsum/sum;
 
     const float d_by_nrows = *grad/gridDim.x;
     for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) {

src/llama-context.cpp

@@ -3483,6 +3483,9 @@ void llama_context::opt_epoch_iter(
                 const float onef = 1.0f;
                 for (uint32_t pos_ubatch = 0; pos_ubatch < n_ubatch; ++pos_ubatch) {
                     const uint32_t ilabel = pos_ctx + pos_batch + pos_ubatch;
+                    if (labels_sparse[ilabel] < 0) {
+                        continue; // masked: carries no loss (the backward honours zero label rows)
+                    }
                     GGML_ASSERT(labels_sparse[ilabel] < labels->ne[0]);
                     ggml_backend_tensor_set(labels, &onef, (pos_ubatch*labels->ne[0] + labels_sparse[ilabel])*sizeof(float), sizeof(float));
                 }

tests/test-backend-ops.cpp

@@ -7363,17 +7363,41 @@ struct test_flash_attn_ext : public test_case {
 };
 
 // GGML_OP_CROSS_ENTROPY_LOSS
+// Completion-only training masks rows: a masked position's label row is all zero. Every other
+// row sums to 1. "masked" zeroes every other row, so the finite-difference check sees whether the
+// backward gives those rows a zero gradient (it must: they carry no loss).
+static ggml_tensor * ce_test_labels(ggml_context * ctx, ggml_tensor * labels, bool masked) {
+    labels = ggml_soft_max(ctx, labels);
+    ggml_set_name(labels, "labels_normalized");
+    if (masked) {
+        ggml_tensor * mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, labels->ne[1], labels->ne[2], labels->ne[3]);
+        ggml_set_name(mask, "mask");
+        labels = ggml_mul(ctx, labels, mask);
+        ggml_set_name(labels, "labels_masked");
+    }
+    return labels;
+}
+
+static void ce_test_init_mask(ggml_tensor * t) {
+    std::vector<float> m(ggml_nelements(t));
+    for (size_t i = 0; i < m.size(); ++i) {
+        m[i] = (i % 2 == 0) ? 1.0f : 0.0f;
+    }
+    ggml_backend_tensor_set(t, m.data(), 0, m.size()*sizeof(float));
+}
+
 struct test_cross_entropy_loss : public test_case {
     const ggml_type type;
     const std::array<int64_t, 4> ne;
+    const bool masked;
 
     std::string vars() override {
-        return VARS_TO_STR2(type, ne);
+        return VARS_TO_STR3(type, ne, masked);
     }
 
     test_cross_entropy_loss(ggml_type type = GGML_TYPE_F32,
-            std::array<int64_t, 4> ne = {10, 5, 4, 3})
-        : type(type), ne(ne) {}
+            std::array<int64_t, 4> ne = {10, 5, 4, 3}, bool masked = false)
+        : type(type), ne(ne), masked(masked) {}
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         ggml_tensor * logits = ggml_new_tensor(ctx, type, 4, ne.data());
@@ -7384,9 +7408,8 @@ struct test_cross_entropy_loss : public test_case {
         // The labels are assumed to be constant -> no gradients.
         ggml_set_name(labels, "labels");
 
-        // Ensure labels add up to 1:
-        labels = ggml_soft_max(ctx, labels);
-        ggml_set_name(labels, "labels_normalized");
+        // Rows add up to 1 (or to 0 where masked):
+        labels = ce_test_labels(ctx, labels, masked);
 
         ggml_tensor * out = ggml_cross_entropy_loss(ctx, logits, labels);
         ggml_set_name(out, "out");
@@ -7397,12 +7420,21 @@ struct test_cross_entropy_loss : public test_case {
     void initialize_tensors(ggml_context * ctx) override {
         // For larger abs. diffs between logits softmax is more linear, therefore more precise num. gradients.
         for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
-            init_tensor_uniform(t, -100.0f, 100.0f);
+            if (strcmp(t->name, "mask") == 0) {
+                ce_test_init_mask(t);
+            } else if (masked) {
+                // No exact zeros anywhere: the MAA metric is (a-b)/(a+b), so two exact zeros make
+                // it NaN and a NaN error reads as a pass. At +-100 softmax underflows to 0 and the
+                // check cannot fail; at +-3 every probability is > 1e-3 and it can.
+                init_tensor_uniform(t, -3.0f, 3.0f);
+            } else {
+                init_tensor_uniform(t, -100.0f, 100.0f);
+            }
         }
     }
 
     float grad_eps() override {
-        return 1.0f;
+        return masked ? 1e-1f : 1.0f;
     }
 
     bool grad_precise() override {
@@ -7414,14 +7446,25 @@ struct test_cross_entropy_loss : public test_case {
 struct test_cross_entropy_loss_back : public test_case {
     const ggml_type type;
     const std::array<int64_t, 4> ne;
+    const bool masked;
 
     std::string vars() override {
-        return VARS_TO_STR2(type, ne);
+        return VARS_TO_STR3(type, ne, masked);
     }
 
     test_cross_entropy_loss_back(ggml_type type = GGML_TYPE_F32,
-            std::array<int64_t, 4> ne = {10, 5, 4, 3})
-        : type(type), ne(ne) {}
+            std::array<int64_t, 4> ne = {10, 5, 4, 3}, bool masked = false)
+        : type(type), ne(ne), masked(masked) {}
+
+    void initialize_tensors(ggml_context * ctx) override {
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+            if (strcmp(t->name, "mask") == 0) {
+                ce_test_init_mask(t);
+            } else {
+                init_tensor_uniform(t);
+            }
+        }
+    }
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         ggml_tensor * grad = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
@@ -7433,9 +7476,8 @@ struct test_cross_entropy_loss_back : public test_case {
         ggml_tensor * labels = ggml_new_tensor(ctx, type, 4, ne.data());
         ggml_set_name(labels, "labels");
 
-        // Ensure labels add up to 1:
-        labels = ggml_soft_max(ctx, labels);
-        ggml_set_name(labels, "labels_normalized");
+        // Rows add up to 1 (or to 0 where masked):
+        labels = ce_test_labels(ctx, labels, masked);
 
         ggml_tensor * out = ggml_cross_entropy_loss_back(ctx, grad, logits, labels);
         ggml_set_name(out, "out");
@@ -10168,6 +10210,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_cross_entropy_loss     (GGML_TYPE_F32, {30000, 1, 1, 1}));
     test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, {   10, 5, 4, 3}));
     test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, {30000, 1, 1, 1}));
+    // masked rows (completion-only training): zero label rows must carry zero gradient
+    test_cases.emplace_back(new test_cross_entropy_loss     (GGML_TYPE_F32, {   10, 5, 4, 3}, true));
+    test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, {   10, 5, 4, 3}, true));
+    test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, {30000, 4, 1, 1}, true));
 
     test_cases.emplace_back(new test_opt_step_adamw(GGML_TYPE_F32, {10, 5, 4, 3}));
     test_cases.emplace_back(new test_opt_step_sgd(GGML_TYPE_F32, {10, 5, 4, 3}));

tools/server/server-context.cpp

@@ -4134,7 +4134,15 @@ bool server_context::load_model(common_params & params) {
     if (ok && impl->model_tgt != nullptr) {
         server_context_impl * p = impl.get();
         impl->trainer = std::make_unique<server_trainer>(impl->model_tgt, impl->params_base,
-                                                         [p]() { return p->n_busy_slots.load(); });
+                                                         [p]() { return p->n_busy_slots.load(); },
+                                                         // the SAME templates serving frames turns with
+                                                         [p](const json & messages, bool add_generation_prompt) {
+                                                             common_chat_templates_inputs in;
+                                                             in.messages              = common_chat_msgs_parse_oaicompat(messages);
+                                                             in.add_generation_prompt = add_generation_prompt;
+                                                             in.use_jinja             = p->params_base.use_jinja;
+                                                             return common_chat_templates_apply(p->chat_params.tmpls.get(), in).prompt;
+                                                         });
     }
     return ok;
 }