Function bertExample

examples/cpp/bert_fp8/bert_fp8_example_squad.cc:57–235

template<typename T1, typename T2>
int bertExample(size_t max_batch_size,
                size_t num_layers,
                size_t head_num,
                size_t size_per_head,
                bool   is_remove_padding,
                bool   allow_gemm_test)
{
    const int fp8_mode = 2;
    FT_LOG_INFO("fp8_mode: %d", fp8_mode);
    FT_LOG_INFO("Device: %s \n", getDeviceName().c_str());

    const size_t hidden_units            = head_num * size_per_head;
    const size_t inter_size              = 4 * hidden_units;
    const size_t vocab_size              = 30522;  // Fixed by bert model config
    const size_t max_position_embeddings = 512;    // Fixed by bert model config
    const size_t token_type_vocab_size   = 2;      // Fixed by bert model config

    cudaStream_t     stream;
    cublasHandle_t   cublas_handle;
    cublasLtHandle_t cublaslt_handle;
    cudaStreamCreate(&stream);
    cublasCreate(&cublas_handle);
    cublasLtCreate(&cublaslt_handle);
#ifdef SPARSITY_ENABLED
    cusparseLtHandle_t cusparselt_handle;
    CHECK_CUSPARSE(cusparseLtInit(&cusparselt_handle));
#endif
    cublasSetStream(cublas_handle, stream);
    cublasAlgoMap* cublas_algo_map = new cublasAlgoMap("gemm_config.in", "");

    Allocator<AllocatorType::CUDA> allocator(getDevice());

    std::mutex* cublas_wrapper_mutex = new std::mutex();
#ifdef SPARSITY_ENABLED
    // cublasFP8MMWrapper cublas_wrapper = cublasFP8MMWrapper(
    //     cublas_handle, cublaslt_handle, cusparselt_handle, stream, cublas_algo_map, cublas_wrapper_mutex,
    //     &allocator);
#else
    cublasFP8MMWrapper cublas_wrapper =
        cublasFP8MMWrapper(cublas_handle, cublaslt_handle, stream, cublas_algo_map, cublas_wrapper_mutex, &allocator);
#endif

    // Prepare FP8 bert model and weight
    BertFP8Weight<T1, T2> bert_weights(hidden_units,
                                       head_num,
                                       size_per_head,
                                       inter_size,
                                       num_layers,
                                       vocab_size,
                                       max_position_embeddings,
                                       token_type_vocab_size,
                                       1,
                                       1,
                                       fp8_mode,
                                       true,
                                       true);
    bert_weights.loadModel("/home/scratch.bhsueh_sw_1/FP8/mlperf/ft_fp8/");
    bert_weights.transposeWeight();