Function bertExample

examples/cpp/bert_fp8/bert_fp8_example.cc:157–452

template<typename T1, typename T2>
int bertExample(size_t       batch_size,
                size_t       num_layers,
                size_t       seq_len,
                size_t       head_num,
                size_t       size_per_head,
                bool         is_remove_padding,
                std::string  model_path,
                const float& mean_rel_diff_threshold)
{
    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 = cublasMMWrapper(
    //     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

    const size_t input_size     = batch_size * seq_len * head_num * size_per_head;
    std::string  fp32_ckpt_path = model_path + "/ft_fp32/1-gpu/";
    std::string  fp8_ckpt_path  = model_path + "/ft_fp8/1-gpu/";
    std::string  input_path     = model_path + "/bert_input_unit_test/bs_";  // use random ids

    // Preare FP16 bert model and weight
    BertWeight<half> f_bert_weights(hidden_units, inter_size, num_layers);
    f_bert_weights.loadModel(fp32_ckpt_path);

    AttentionType fp8_attention_type = getAttentionType<T1>(size_per_head, getSMVersion(), is_remove_padding, seq_len);
    AttentionType f_attention_type = getAttentionType<half>(size_per_head, getSMVersion(), is_remove_padding, seq_len);
    cublasMMWrapper f_cublas_wrapper =
        cublasMMWrapper(cublas_handle, cublaslt_handle, stream, cublas_algo_map, cublas_wrapper_mutex, &allocator);
    f_cublas_wrapper.setFP16GemmConfig();