Source browser
examples/transformer.cpp
Read the complete file below or download its original bytes. Line numbers provide stable links to a declaration or processing step.
File fingerprint
1/** @file transformer.cpp
2 * @brief Train and validate a tiny transformer on the CPU or selected GL backend.
3 * Analytic forward derivatives train every parameter. Finite differences are
4 * used only to validate the gradient implementation before optimization.
5 */
6#include "runner.hpp" // Select CPU execution or the requested graphics profile.
7#include "transformer_support.hpp" // Use the reusable tiny transformer model.
8#include "transformer_reference.hpp" // Use an independent double-precision prediction oracle.
9#include <algorithm> // Use standard copying, limits and zero filling.
10#include <cmath> // Use ordinary scalar math in host designers and oracles.
11#include <iostream> // Report verified results and useful failures.
12#include <vector> // Use owned host arrays for coefficients and reference results.
13
14namespace { // Keep these example helpers in their own namespace.
15
16void check_gradients(shdr_examples::transformer_model& model, shdr_examples::runner& run) { // Compare analytic derivatives with the independent double oracle.
17 run.run(model.gradient_pipeline()); // Execute the linked stages on the selected backend.
18 const float* derivative = static_cast<const shdr::buffer<float>&>(model.gradients).data(); // Read the resource after any required GPU result synchronization.
19 std::vector<float> analytic(model.parameter_count, 0.0f); // Keep the averaged analytic derivative for every parameter.
20 for (unsigned parameter = 0; parameter < model.parameter_count; ++parameter) // Visit every trainable parameter.
21 for (unsigned row = 0; row < model.batch_count(); ++row) // Visit each token sequence in the batch.
22 analytic[parameter] += // Average this analytic derivative over the batch.
23 derivative[row * model.parameter_count + parameter] / model.batch_count(); // Read the sample-major gradient for this parameter.
24 const float* original = static_cast<const shdr::buffer<float>&>(model.parameters).data(); // Read the resource after any required GPU result synchronization.
25 std::vector<double> parameters(original, original + model.parameter_count); // Copy model parameters into the independent double reference.
26 const float* inputs = static_cast<const shdr::buffer<float>&>(model.inputs).data(); // Read the resource after any required GPU result synchronization.
27 const unsigned* labels = static_cast<const shdr::buffer<unsigned>&>(model.labels).data(); // Read the resource after any required GPU result synchronization.
28 run.run(model.prediction_pipeline()); // Execute the linked stages on the selected backend.
29 const shdr::glsl::vec4* predictions = // Read the resource after any required GPU result synchronization.
30 static_cast<const shdr::buffer<shdr::glsl::vec4>&>(model.predictions).data(); // Read probabilities, losses and classes for the double-oracle check.
31 for (unsigned row = 0; row < model.batch_count(); ++row) { // Visit each token sequence in the batch.
32 const std::array<double, 4> expected = shdr_examples::detail::transformer_reference( // Compute the independent double probabilities and loss for this sequence.
33 inputs + 16u * row, parameters.data(), labels[row]); // Use this sequence and its label in the double reference.
34 for (unsigned c = 0; c < 4; ++c) // Visit every value needed by this calculation.
35 shdr_examples::require_close(predictions[row][c], // Compare this result with its independent expected value.
36 static_cast<float>(expected[c]), // Convert the computed host value to the shader float type.
37 "transformer double prediction oracle", // Identify this check clearly in any failure message.
38 0.00002f); // Apply the declared float error allowance.
39 } // Finish this channel, sample or feature pass.
40 const double step = 0.00001; // Use a small perturbation only for derivative validation.
41 for (unsigned parameter = 0; parameter < model.parameter_count; ++parameter) { // Visit every trainable parameter.
42 const double value = parameters[parameter]; // Save this double parameter before perturbing it.
43 double above = 0.0, below = 0.0; // Prepare the two perturbed-reference loss accumulators.
44 parameters[parameter] = value + step; // Perturb this double reference parameter upward.
45 for (unsigned row = 0; row < model.batch_count(); ++row) // Visit each token sequence in the batch.
46 above += shdr_examples::detail::transformer_reference( // Accumulate the independent perturbed-reference loss.
47 inputs + 16u * row, parameters.data(), labels[row])[2]; // Use this sequence and its label in the double reference.
48 parameters[parameter] = value - step; // Perturb this double reference parameter downward.
49 for (unsigned row = 0; row < model.batch_count(); ++row) // Visit each token sequence in the batch.
50 below += shdr_examples::detail::transformer_reference( // Accumulate the independent perturbed-reference loss.
51 inputs + 16u * row, parameters.data(), labels[row])[2]; // Use this sequence and its label in the double reference.
52 parameters[parameter] = value; // Restore the double reference parameter after its check.
53 const float numerical = // Compute the central-difference mean-loss derivative.
54 static_cast<float>((above - below) / (2.0 * step * model.batch_count())); // Convert the computed host value to the shader float type.
55 const float tolerance = // Allow float rounding while keeping the derivative check strict.
56 0.00002f + 0.001f * std::max(std::fabs(numerical), std::fabs(analytic[parameter])); // Combine a small absolute error allowance with a relative one.
57 shdr_examples::require_close(analytic[parameter], // Compare this result with its independent expected value.
58 numerical, // Supply the independently computed expected value.
59 "transformer analytic gradient " + std::to_string(parameter), // Identify this check clearly in any failure message.
60 tolerance); // Apply the declared float error allowance.
61 } // Finish this channel, sample or feature pass.
62} // Finish the check_gradients operation.
63
64void check_trained_cpu_parity(shdr_examples::transformer_model& gpu) { // Compare every trained CPU and GPU parameter and prediction.
65 shdr_examples::transformer_model cpu(gpu.batch_count()); // Build an independent CPU model with the same seed.
66 for (unsigned step = 0; step < 32u; ++step) // Visit every value needed by this calculation.
67 cpu.training_pipeline().run(); // Apply the same analytic update independently on the CPU.
68 cpu.prediction_pipeline().run(); // Compute the independently trained CPU predictions.
69 const float* expected = static_cast<const shdr::buffer<float>&>(cpu.parameters).data(); // Read the resource after any required GPU result synchronization.
70 const float* actual = static_cast<const shdr::buffer<float>&>(gpu.parameters).data(); // Read the resource after any required GPU result synchronization.
71 for (unsigned i = 0; i < cpu.parameter_count; ++i) // Visit every trainable parameter.
72 shdr_examples::require_close( // Compare this result with its independent expected value.
73 actual[i], expected[i], "trained transformer CPU/GPU parameter", 0.0003f); // Supply the actual and expected result values for comparison.
74 expected = static_cast<const shdr::buffer<float>&>(cpu.gradients).data(); // Read the corresponding CPU and GPU result resource.
75 actual = static_cast<const shdr::buffer<float>&>(gpu.gradients).data(); // Read the corresponding CPU and GPU result resource.
76 for (unsigned i = 0; i < cpu.gradients.size(); ++i) // Visit every value needed by this calculation.
77 shdr_examples::require_close( // Compare this result with its independent expected value.
78 actual[i], expected[i], "trained transformer CPU/GPU gradient", 0.0001f); // Supply the actual and expected result values for comparison.
79 const shdr::glsl::vec4* expected_predictions = // Read the resource after any required GPU result synchronization.
80 static_cast<const shdr::buffer<shdr::glsl::vec4>&>(cpu.predictions).data(); // Read the independently trained CPU prediction records.
81 const shdr::glsl::vec4* actual_predictions = // Read the resource after any required GPU result synchronization.
82 static_cast<const shdr::buffer<shdr::glsl::vec4>&>(gpu.predictions).data(); // Read the trained GPU prediction records.
83 for (unsigned row = 0; row < cpu.batch_count(); ++row) // Visit each token sequence in the batch.
84 for (unsigned c = 0; c < 4; ++c) // Visit every value needed by this calculation.
85 shdr_examples::require_close(actual_predictions[row][c], // Compare this result with its independent expected value.
86 expected_predictions[row][c], // Supply the independently computed expected value.
87 "trained transformer CPU/GPU prediction", // Identify this check clearly in any failure message.
88 0.0001f); // Apply the declared float error allowance.
89} // Finish the check_trained_cpu_parity operation.
90
91void train(shdr_examples::runner& run) { // Verify the complete transformer training example.
92 shdr_examples::transformer_model model(8); // Allocate the seeded eight-sequence transformer fixture.
93 check_gradients(model, run); // Run the corresponding independent correctness check.
94 model.reset_parameters(); // Restore the seed before measuring complete training convergence.
95 run.run(model.prediction_pipeline()); // Execute the linked stages on the selected backend.
96 const float initial = model.loss(); // Keep the loss before the first training update.
97 const float* start = static_cast<const shdr::buffer<float>&>(model.parameters).data(); // Read the resource after any required GPU result synchronization.
98 const std::vector<float> original(start, start + model.parameter_count); // Keep initial parameters for the learned-family update check.
99 for (unsigned step = 0; step < 32u; ++step) // Visit every value needed by this calculation.
100 run.run(model.training_pipeline()); // Execute the linked stages on the selected backend.
101 run.run(model.prediction_pipeline()); // Execute the linked stages on the selected backend.
102 const float final = model.loss(); // Read the loss after all analytic training steps.
103 const unsigned correct = model.correct(); // Count the trained model predictions matching their labels.
104 if (!std::isfinite(final) || final >= initial * 0.35f || correct != model.batch_count()) // Reject a nonfinite or out-of-range value.
105 throw shdr::error("transformer training did not converge on the seeded dataset"); // Report the invalid design or failed check clearly.
106 // Each learned family must participate in optimization, including Q and K.
107 const unsigned boundaries[] = {0, 16, 32, 48, 64, 80, 84, 100, 104, 120, 124, 132, 134}; // Mark the learned parameter families so every family is checked.
108 const float* trained = static_cast<const shdr::buffer<float>&>(model.parameters).data(); // Read the resource after any required GPU result synchronization.
109 for (unsigned family = 0; family + 1u < sizeof(boundaries) / sizeof(boundaries[0]); ++family) { // Visit every value needed by this calculation.
110 bool changed = false; // Start checking this learned parameter family.
111 for (unsigned parameter = boundaries[family]; parameter < boundaries[family + 1u]; // Visit every trainable parameter.
112 ++parameter) // Advance through every parameter in this learned family.
113 changed = changed || trained[parameter] != original[parameter]; // Notice whether any parameter in this family was updated.
114 if (!changed) // Require this learned parameter family to participate.
115 throw shdr::error("a transformer parameter family did not train"); // Report the invalid design or failed check clearly.
116 } // Finish this channel, sample or feature pass.
117 if (run.gpu()) // Use the requested GPU path or explicit CPU transition.
118 check_trained_cpu_parity(model); // Run the corresponding independent correctness check.
119 std::cout << "Transformer: 134 analytic parameter gradients verified. Cross entropy " << initial // Print the verified result or the reason a check failed.
120 << " -> " << final << ", classification " << correct << '/' << model.batch_count() // Complete the readable result message.
121 << " after 32 SGD steps\n"; // Complete the readable result message.
122} // Finish the train operation.
123} // namespace
124
125int main(int argc, char** argv) { // Run the selected CPU or GPU example.
126 try { // Run the operation while preserving its failure diagnostic.
127 shdr_examples::runner run(argc, argv); // Select CPU execution or the requested GPU profile.
128 train(run); // Run the complete transformer gradient and training checks.
129 return 0; // Report successful validation to the caller.
130 } // Finish the checked operation.
131#ifdef SHDR_HAVE_EGL // Enable this block only for the indicated compilation environment.
132 catch (const shdr_examples::context_unavailable& e) { // Handle this expected backend or validation error.
133 std::cerr << "SKIP: " << e.what() << '\n'; // Print the verified result or the reason a check failed.
134 return 77; // Mark the unavailable graphics profile as skipped.
135 } // Finish handling this reported failure.
136#endif // Finish the conditional compilation block.
137 catch (const std::exception& e) { // Handle this expected backend or validation error.
138 std::cerr << e.what() << '\n'; // Print the verified result or the reason a check failed.
139 return 1; // Report the validation failure to the caller.
140 } // Finish handling this reported failure.
141} // Finish the main operation.