diff --git a/dev/cuda/README.md b/dev/cuda/README.md index d020cf6..22ad4d0 100644 --- a/dev/cuda/README.md +++ b/dev/cuda/README.md @@ -7,7 +7,7 @@ See the top of each file for how to compile and run the kernel. Alternatively, t For example, we can look at the top of `layernorm_forward.cu` to build the forward pass kernels for the LayerNorm: ```bash -nvcc -O3 --use_fast_math layernorm_forward.cu -o layernorm_forward +nvcc -O3 --use_fast_math -lcublas -lcublasLt layernorm_forward.cu -o layernorm_forward ``` or simply diff --git a/dev/cuda/classifier_fused.cu b/dev/cuda/classifier_fused.cu index 2125b87..5b6c986 100644 --- a/dev/cuda/classifier_fused.cu +++ b/dev/cuda/classifier_fused.cu @@ -664,7 +664,7 @@ void fused_classifier5(float* dlogits, float* losses, int B, int T, int V, int P, int block_size) { const int N = B * T; const int grid_size = N; - fused_classifier_kernel5<<>>((floatX*)dlogits, (floatX*)losses, NULL, (floatX*)logits, (floatX*)dlosses, targets, B, T, V, P); + fused_classifier_kernel5<<>>((floatX*)dlogits, (floatX*)losses, NULL, (floatX*)logits, (floatX*)dlosses, targets, B, T, V, P); cudaCheck(cudaGetLastError()); } diff --git a/dev/cuda/encoder_backward.cu b/dev/cuda/encoder_backward.cu index 5322187..7c14d0e 100644 --- a/dev/cuda/encoder_backward.cu +++ b/dev/cuda/encoder_backward.cu @@ -163,14 +163,17 @@ int main(int argc, char **argv) { } printf("Using kernel %d\n", kernel_num); - // set up block sizes + // first check the correctness of the kernel + encoder_backward_cpu(dwte, dwpe, dout, inp, B, T, C); + + // time the kernel at different block sizes int block_sizes[] = {32, 64, 128, 256, 512, 1024}; - // first check the correctness of the kernel for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) { int block_size = block_sizes[j]; + cudaCheck(cudaMemset(d_dwte, 0, V * C * sizeof(float))); + cudaCheck(cudaMemset(d_dwpe, 0, T * C * sizeof(float))); printf("Checking block size %d.\n", block_size); - encoder_backward_cpu(dwte, dwpe, dout, inp, B, T, C); encoder_backward(kernel_num, d_dwte, d_dwpe, d_dout, d_inp, B, T, C, block_size); validate_result(d_dwte, dwte, "dwte", V * C, 1e-5f); validate_result(d_dwpe, dwpe, "dwpe", T * C, 1e-5f); diff --git a/llmc/mfu.h b/llmc/mfu.h new file mode 100644 index 0000000..ec26537 --- /dev/null +++ b/llmc/mfu.h @@ -0,0 +1,136 @@ +#ifndef MFU_H +#define MFU_H + +#include +#include +#include + +// tied to enum PrecisionMode, in a future refactor make them the same +#define MFUH_PRECISION_FP32 0 +#define MFUH_PRECISION_FP16 1 +#define MFUH_PRECISION_BF16 2 + +typedef struct { + float TF_32; // tensor-core performance 32 bit + float BF_16_32; // bf16 with 32 bit accumulate + float FP_16_32; // fp16 with 32 bit accumulate + float FP_16_16; // fp16 with 16 bit accumulate + float FP_8_32; // and so on + float FP_8_16; + float CLOCK; // clock frequency from the spec sheet + float CORES; // #TCs from the spec sheet +} PerfData; + +// basic default data from the nvidia whitepapers +static const PerfData VOLTA = {125.0f, -1.f, 125.f, -1.f, -1.f, -1.f, 1530.f, 640.f}; +static const PerfData AMPERE_DATACENTER = {156.f, 312.f, 312.f, 312.f, -1.f, -1.f, 1410.f, 432.f}; +static const PerfData AMPERE_CONSUMER = {40.f, 80.f, 80.f, 160.f, -1.f, -1.f, 1860.f, 336.f}; +static const PerfData HOPPER = {378.f, 756.f, 756.f, 756.f, 1513.f, 1513.f, 1620.f, 456.f}; +static const PerfData ADA = {82.6f, 165.2f, 165.2f, 330.3f, 330.3f, 660.6f, 2520.f, 512.f}; + +typedef struct { + const char* name; + const PerfData* perf_data; + float new_cores; + float new_mhz; +} GPUEntry; + +// the overrides for each specific GPU +static GPUEntry gpu_db[] = { + {"Tesla V100-SXM2-16GB", &VOLTA, 640, 1530}, + {"Tesla V100-PCIE-32GB", &VOLTA, 640, 1530}, + {"NVIDIA A100-PCIE-40GB", &ERE_DATACENTER, 432, 1410}, + {"NVIDIA A100-PCIE-80GB", &ERE_DATACENTER, 432, 1410}, + {"NVIDIA A100-SXM4-40GB", &ERE_DATACENTER, 432, 1410}, + {"NVIDIA A100-SXM4-80GB", &ERE_DATACENTER, 432, 1410}, + {"NVIDIA RTX A2000", &ERE_CONSUMER, 104, 1200}, + {"NVIDIA RTX A4000", &ERE_CONSUMER, 192, 1560}, + {"NVIDIA RTX A4500", &ERE_CONSUMER, 224, 1650}, + {"NVIDIA RTX A5000", &ERE_CONSUMER, 256, 1695}, + {"NVIDIA RTX A5500", &ERE_CONSUMER, 320, 1770}, + {"NVIDIA RTX A6000", &ERE_CONSUMER, 336, 1800}, + {"NVIDIA GeForce RTX 3090 Ti", &ERE_CONSUMER, 336, 1860}, + {"NVIDIA GeForce RTX 3090", &ERE_CONSUMER, 328, 1695}, + {"NVIDIA GeForce RTX 3080 Ti", &ERE_CONSUMER, 320, 1665}, + {"NVIDIA GeForce RTX 3080", &ERE_CONSUMER, 272, 1710}, + {"NVIDIA GeForce RTX 3070 Ti", &ERE_CONSUMER, 192, 1770}, + {"NVIDIA GeForce RTX 3070", &ERE_CONSUMER, 184, 1725}, + {"NVIDIA GeForce RTX 3060 Ti", &ERE_CONSUMER, 152, 1665}, + {"NVIDIA GeForce RTX 3060", &ERE_CONSUMER, 112, 1777}, + {"NVIDIA RTX A2000 ADA", &ADA, 88, 2130}, + {"NVIDIA RTX A4000 ADA", &ADA, 192, 2175}, + {"NVIDIA RTX A4500 ADA", &ADA, 224, 2580}, + {"NVIDIA RTX A5000 ADA", &ADA, 400, 2550}, + {"NVIDIA RTX A5880 ADA", &ADA, 440, 2460}, + {"NVIDIA RTX A6000 ADA", &ADA, 568, 2505}, + {"NVIDIA GeForce RTX 4090", &ADA, 512, 2520}, + {"NVIDIA GeForce RTX 4080 SUPER", &ADA, 320, 2550}, + {"NVIDIA GeForce RTX 4080", &ADA, 304, 2505}, + {"NVIDIA GeForce RTX 4070 Ti SUPER", &ADA, 264, 2610}, + {"NVIDIA GeForce RTX 4070 Ti", &ADA, 240, 2610}, + {"NVIDIA GeForce RTX 4070 SUPER", &ADA, 224, 2475}, + {"NVIDIA GeForce RTX 4070", &ADA, 184, 2475}, + {"NVIDIA GeForce RTX 4070", &ADA, 184, 2475}, + {"NVIDIA GeForce RTX 4060 Ti", &ADA, 136, 2535}, + {"NVIDIA GeForce RTX 4060", &ADA, 96, 2460}, + {"NVIDIA H100 80GB HBM3", &HOPPER, 528, 1830}, // HBM3 = SXM5 +}; + +float get_flops_promised(const char* device, int precision_mode) { + /* + This function is used to estimate the Model Flops Utilization (MFU) + basically we have to figure out how many flops the GPU can do per second. + Note that this is not a simple endeavor and may well go wrong! The details are tricky. + The returned value is in units of 1e12. + + For the non-top models, actual performance numbers aren't that easy to find, e.g., + here https://www.techpowerup.com/gpu-specs/rtx-a4000.c3756, does "Theoretical Performance" + seems to be without tensor cores. + + So, instead we use that all these cards just use the same types of tensor cores in different + numbers and at different frequencies. Then we just need to look up these two easily accesible + numbers for all the other GPUs. + linear scaling seems to work: comparing spec sheet and calculation: + 4080: 304TCs, 2505 GHz; 97.5TFlops = 165.2/512*304 /2520 * 2505 + + Original numbers for the top GPUS are from. + https://resources.nvidia.com/en-us-tensor-core + https://images.nvidia.com/aem-dam/Solutions/geforce/ada/nvidia-ada-gpu-architecture.pdf + */ + + // validate the precision mode as one of the three possible values + if (!(precision_mode == MFUH_PRECISION_FP32 || precision_mode == MFUH_PRECISION_FP16 || precision_mode == MFUH_PRECISION_BF16)) { + fprintf(stderr, "Invalid precision mode: %d\n", precision_mode); + return -1.0f; + } + + // do a linear search until you find our GPU, then calculate the flops promised + int num_gpu_entries = sizeof(gpu_db) / sizeof(gpu_db[0]); + for (int i = 0; i < num_gpu_entries; i++) { + if (strcmp(gpu_db[i].name, device) == 0) { + const PerfData* perf_data = gpu_db[i].perf_data; + + // look up the default flops value for the given precision mode + float value = -1.0f; + if (precision_mode == MFUH_PRECISION_BF16) { value = perf_data->BF_16_32; } + if (precision_mode == MFUH_PRECISION_FP32) { value = perf_data->TF_32; } + if (precision_mode == MFUH_PRECISION_FP16) { value = perf_data->FP_16_32; } + + // we'd get here if we're e.g. trying to use BF16 on Volta GPU or something... + if (value < 0.0f) { + fprintf(stderr, "No data for GPU %s and precision mode %d\n", device, precision_mode); + return -1.0f; + } + + // adjust flops based on the specific core count and clock frequency of this GPU + float new_cores = gpu_db[i].new_cores; + float new_mhz = gpu_db[i].new_mhz; + float adjusted = value * (new_cores / perf_data->CORES) * (new_mhz / perf_data->CLOCK); + return adjusted; + } + } + + return -1.0f; // ¯\_(ツ)_/¯ +} + +#endif // MFU_H diff --git a/llmc/rand.h b/llmc/rand.h index ba13de9..60ed393 100644 --- a/llmc/rand.h +++ b/llmc/rand.h @@ -163,8 +163,8 @@ void uniform_(float* data, unsigned int numel, float from, float to, mt19937_sta } } -// Box�Muller transform - +// Box-Muller transform: maps uniform random numbers to Gaussian distributed numbers +// https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform void normal_fill_16(float* data, float mean, float std, mt19937_state* state) { #define EPSILONE 1e-12 for (unsigned int t = 0; t < 8; t++) { diff --git a/train_gpt2.cu b/train_gpt2.cu index 143750e..06da2f5 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -7,6 +7,8 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage. #include #include #include +#include +#include #include #include #include @@ -38,6 +40,8 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage. #include "llmc/sampler.h" // defines: logger_init, logger_log_eval, logger_log_val, logger_log_train #include "llmc/logger.h" +// defines: get_flops_promised +#include "llmc/mfu.h" // ---------------------------------------------------------------------------- // CUDA precision settings @@ -228,10 +232,10 @@ struct alignas(16) Packed128 { return result; } __device__ static Packed128 zeros() { - return constant(0); + return constant(0.f); } __device__ static Packed128 ones() { - return constant(1); + return constant(1.f); } __device__ ElementType& operator[](int index) { @@ -942,12 +946,12 @@ __global__ void gelu_forward_kernel2(floatX* out, const floatX* inp) { store128(out + idx, packed_out); } -__global__ void gelu_backward_kernel(floatX* dinp, const floatX* inp, const floatX* dout) { +__global__ void gelu_backward_inplace_kernel(floatX* d_in_out, const floatX* inp) { int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; x128 packed_dinp; x128 packed_inp = load128cs(inp + idx); - x128 packed_dout = load128cs(dout + idx); + x128 packed_dout = load128(d_in_out + idx); for (int k = 0; k < packed_inp.size; ++k) { float x = (float)packed_inp[k]; float cube = 0.044715f * x * x * x; @@ -958,7 +962,7 @@ __global__ void gelu_backward_kernel(floatX* dinp, const floatX* inp, const floa float local_grad = 0.5f * (1.0f + tanh_out) + x * 0.5f * sech_out * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x); packed_dinp[k] = (floatX)(local_grad * (float)packed_dout[k]); } - store128(dinp + idx, packed_dinp); + store128(d_in_out + idx, packed_dinp); } template @@ -1747,12 +1751,12 @@ void gelu_forward(floatX* out, const floatX* inp, int N) { cudaCheck(cudaGetLastError()); } -void gelu_backward(floatX* dinp, const floatX* inp, const floatX* dout, const int N) { +void gelu_backward_inplace(floatX* d_in_out, const floatX* inp, const int N) { NVTX_RANGE_FN(); const int block_size = 128; assert(N % block_size == 0); const int grid_size = CEIL_DIV(N, block_size * x128::size); - gelu_backward_kernel<<>>(dinp, inp, dout); + gelu_backward_inplace_kernel<<>>(d_in_out, inp); cudaCheck(cudaGetLastError()); } @@ -1874,7 +1878,7 @@ void fused_classifier(Type* logits, Type* losses, const int block_size = 1024; const int N = B * T; const int grid_size = N; - fused_classifier_kernel5<<>>(logits, losses, (floatX*)NULL, dloss, targets, B, T, V, P); + fused_classifier_kernel5<<>>(logits, losses, (floatX*)NULL, dloss, targets, B, T, V, P); cudaCheck(cudaGetLastError()); } @@ -2147,13 +2151,43 @@ typedef struct { floatX* cpu_losses; // CPU buffer to copy the losses to, allocated with cudaMallocHost float* cpu_losses_fp32; // same but fp32 unsigned long long rng_state; // the RNG state for seeding stochastic rounding etc. - int use_master_weights; - int recompute; + int use_master_weights; // keep master weights copy in float for optim update? 0|1 + int recompute; // recompute gelu | layernorm forward during model backward? 0|1|2 // todo - if other functions need cpu scratch buffers in the future, reuse as generic scratch? int* workload_indices; // encoder_backward, B*T*num_c_groups (int) int4* bucket_info; // encoder_backward, B*T*num_c_groups (int4) - size for worst case } GPT2; +void gpt2_init_common(GPT2 *model) { + // common inits outside of the model weights + // the weights are initialized either in: + // - gpt2_build_from_checkpoint() if loading from a checkpoint + // - gpt2_build_from_random() if starting from scratch + // memory lazily initialized in forward() + model->acts_memory = NULL; + model->inputs = NULL; + model->targets = NULL; + model->cpu_losses = NULL; + model->cpu_losses_fp32 = NULL; + // the B,T params are determined and set, fixed on first batch in forward() + model->batch_size = 0; + model->seq_len = 0; + model->mean_loss = -1.0f; // -1.0f designates no loss, set at end of forward() + // memory lazily initialized in backward() + model->grads_memory = NULL; + model->grads_acts_memory = NULL; + model->workload_indices = NULL; // on cpu, for encoder_backward + model->bucket_info = NULL; // on cpu, for encoder_backward + // memory lazily initialized in update() + model->m_memory = NULL; + model->v_memory = NULL; + model->master_weights = NULL; + // other default settings + model->rng_state = 13371337; // used in stochastic rounding + model->use_master_weights = 1; // safe default: do keep master weights in fp32 + model->recompute = 1; // good default: recompute gelu but not layernorm +} + void gpt2_write_to_checkpoint(GPT2 *model, const char* checkpoint_path) { // write the model to a checkpoint file printf0("Writing model to %s\n", checkpoint_path); @@ -2243,25 +2277,7 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { free(params_memory_cpu); fcloseCheck(model_file); - // other inits - model->acts_memory = NULL; - model->grads_memory = NULL; - model->m_memory = NULL; - model->v_memory = NULL; - model->master_weights = NULL; - model->grads_acts_memory = NULL; - model->inputs = NULL; - model->targets = NULL; - model->cpu_losses = NULL; - model->cpu_losses_fp32 = NULL; - model->workload_indices = NULL; - model->bucket_info = NULL; - model->batch_size = 0; - model->seq_len = 0; - model->mean_loss = -1.0f; // -1.0f will designate no loss - model->rng_state = 13371337; - model->use_master_weights = 1; // keep master weights copy in float for optim update? - model->recompute = 1; // default to recompute gelu during backward + gpt2_init_common(model); } void gpt2_build_from_random(GPT2 *model, int depth) { @@ -2350,23 +2366,7 @@ void gpt2_build_from_random(GPT2 *model, int depth) { cudaCheck(cudaMemcpy(model->params_memory, params_memory_cpu, model->num_parameters_bytes, cudaMemcpyHostToDevice)); free(params_memory_cpu); - // other inits and defaults - model->acts_memory = NULL; - model->grads_memory = NULL; - model->m_memory = NULL; - model->v_memory = NULL; - model->master_weights = NULL; - model->grads_acts_memory = NULL; - model->inputs = NULL; - model->targets = NULL; - model->cpu_losses = NULL; - model->cpu_losses_fp32 = NULL; - model->batch_size = 0; - model->seq_len = 0; - model->mean_loss = -1.0f; // -1.0f designates no loss - model->rng_state = 13371337; - model->use_master_weights = 1; // keep master weights copy in float for optim update? - model->recompute = 1; // default to recompute gelu during backward + gpt2_init_common(model); } void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, int grad_accum_steps=1) { @@ -2554,7 +2554,7 @@ void gpt2_backward(GPT2 *model, int* inputs) { model->grads_memory = malloc_and_point_parameters(&model->grads, model->param_elements, model->param_sizeof); // we're going to be clever for the activations backward pass. we don't need to exactly // mirror the forward pass activations and we will save memory. - size_t bw_act_sizes[NUM_ACTIVATION_TENSORS]; + size_t bw_act_sizes[NUM_BACKWARD_TENSORS]; fill_in_grad_act_sizes(bw_act_sizes, model->batch_size, model->seq_len, model->config); // count up and allocate the space model->num_grad_acts = 0; @@ -2660,7 +2660,7 @@ void gpt2_backward(GPT2 *model, int* inputs) { gelu_forward(l_fch_gelu, l_fch, B*T*4*C); } matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, scratchF, B, T, 4*C, C); - gelu_backward(dl_bt4c, l_fch, dl_bt4c, B*T*4*C); + gelu_backward_inplace(dl_bt4c, l_fch, B*T*4*C); if(model->recompute >= 2) { // same as gelu above, l_ln1 and l_ln2 are just buffers if recompute >= 2, recompute them here on demand layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C); @@ -2763,10 +2763,24 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // gradient clipping // repurposing this buffer (which isn't needed now) to write grad norm into it float* grad_norm_squared = (float*)model->acts.output; - global_norm_squared(grad_norm_squared, (floatX*)model->grads_memory, model->num_parameters); + if (multi_gpu_config->zero_stage == 1) { + // ^1 because of the ncclReduceScatter() in gpt2_multi_gpu_accumulate, + // grads_memory only contains the averaged gradients at the local shard + // so we only calculate the grad norm at the grads_memory belonging to the local shard + global_norm_squared(grad_norm_squared, grads_memory + shard_offset, shard_num_parameters); + } else { + // the ncclAllReduce() in gpt2_multi_gpu_accumulate has averaged the gradients across all GPUs + // so each GPU can compute the squared norm over the whole grad vector, with no added comms needed + global_norm_squared(grad_norm_squared, grads_memory, model->num_parameters); + } // transfer the gradient norm to CPU float grad_norm_squared_cpu = 0.0f; cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost)); + if (multi_gpu_config->zero_stage == 1) { + // further sum the (partial) squared norm across all GPUs (see comment ^1 above) + grad_norm_squared_cpu = multi_gpu_cpu_float_sum(grad_norm_squared_cpu); + } + if(!isfinite(grad_norm_squared_cpu)) { // may happen due to some issue (e.g. overflow?) // TODO: later may want to keep a global counter of instabilities like this @@ -2830,7 +2844,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // - the position embeddings actively participate at every forward/backward pass float wd = (i == 0 || i == 1 || i == 4 || i == 6 || i == 10 || i == 12) ? weight_decay : 0.0f; // ok finally call the kernel - size_t num_blocks = CEIL_DIV(num_parameters, block_size); + size_t num_blocks = CEIL_DIV(local_params, block_size); adamw_kernel3<<>>(params_ptr, master_ptr, grad_ptr, m_ptr, v_ptr, local_params, learning_rate, beta1, beta2, beta1_correction, beta2_correction, @@ -2870,7 +2884,10 @@ float gpt2_estimate_mfu(GPT2 *model, int num_tokens, float dt) { size_t flops_per_step = flops_per_token * num_tokens; // express our flops throughput as ratio of A100 bfloat16 peak flops float flops_achieved = (float)flops_per_step * (1.0f / dt); // per second - float flops_promised = 312e12f; // A100 GPU bfloat16 peak flops is 312 TFLOPS + float flops_promised = get_flops_promised(deviceProp.name, PRECISION_MODE) * 1e12f; + if(flops_promised < 0) { + return -1.f; // don't know + } float mfu = flops_achieved / flops_promised; return mfu; } @@ -2885,8 +2902,8 @@ void gpt2_free(GPT2 *model) { cudaCheck(cudaFree(model->grads_acts_memory)); cudaCheck(cudaFree(model->inputs)); cudaCheck(cudaFree(model->targets)); - cudaFreeHost(model->cpu_losses); - cudaFreeHost(model->cpu_losses_fp32); + cudaCheck(cudaFreeHost(model->cpu_losses)); + cudaCheck(cudaFreeHost(model->cpu_losses_fp32)); free(model->workload_indices); free(model->bucket_info); } @@ -2985,6 +3002,8 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename cudaCheck(cudaMemcpy(model->m_memory, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); cudaCheck(cudaMemcpy(model->v_memory, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); + free(cpu_buffer); + fclose(state_file); } // ---------------------------------------------------------------------------- @@ -3140,6 +3159,7 @@ int main(int argc, char *argv[]) { ? (cublas_compute == CUBLAS_COMPUTE_32F_FAST_TF32 ? "TF32" : "FP32") : (PRECISION_MODE == PRECISION_FP16 ? "FP16" : "BF16"); printf0("| device | %-50s |\n", deviceProp.name); + printf0("| TFlops | %-50.1f |\n", get_flops_promised(deviceProp.name, PRECISION_MODE)); printf0("| precision | %-50s |\n", precision_str); printf0("+-----------------------+----------------------------------------------------+\n"); @@ -3445,7 +3465,7 @@ int main(int argc, char *argv[]) { } float accumulated_loss = multi_gpu_config.num_processes == 1 ? model.mean_loss : model.accumulated_mean_loss; float mfu = gpt2_estimate_mfu(&model, B * T * grad_accum_steps, time_elapsed_ms / 1000.0f); - printf0("step %4d/%d | train loss %7.6f | norm %6.4f | lr %.2e | %.2f ms | %.1f%% A100 fp16 MFU | %.0f tok/s\n", + printf0("step %4d/%d | train loss %7.6f | norm %6.4f | lr %.2e | %.2f ms | %.1f%% bf16 MFU | %.0f tok/s\n", step + 1, train_num_batches, accumulated_loss, grad_norm, step_learning_rate, time_elapsed_ms, 100*mfu, bias_corrected_ema_tokens_per_second); logger_log_train(&logger, step, model.mean_loss);