Function gptneox_example

examples/cpp/gptneox/gptneox_example.cc:72–504

template<typename T>
void gptneox_example(const INIReader reader)
{
    const std::string model_name = reader.Get("ft_instance_hyperparameter", "model_name");
    std::string       model_dir  = std::string(reader.Get("ft_instance_hyperparameter", "model_dir"));

    int tensor_para_size   = reader.GetInteger("ft_instance_hyperparameter", "tensor_para_size");
    int pipeline_para_size = reader.GetInteger("ft_instance_hyperparameter", "pipeline_para_size");

    const size_t head_num             = reader.GetInteger(model_name, "head_num");
    const size_t size_per_head        = reader.GetInteger(model_name, "size_per_head");
    const size_t vocab_size           = reader.GetInteger(model_name, "vocab_size");
    const size_t decoder_layers       = reader.GetInteger(model_name, "decoder_layers");
    const size_t rotary_embedding_dim = reader.GetInteger(model_name, "rotary_embedding");
    const int    start_id             = reader.GetInteger(model_name, "start_id");
    const int    end_id               = reader.GetInteger(model_name, "end_id");

    const size_t hidden_units = head_num * size_per_head;
    const size_t inter_size   = 4 * hidden_units;

    const size_t beam_width                 = reader.GetInteger("request", "beam_width");
    const uint   top_k                      = (uint)reader.GetInteger("request", "top_k");
    const float  top_p                      = reader.GetFloat("request", "top_p");
    const float  temperature                = reader.GetFloat("request", "temperature");
    const float  repetition_penalty         = reader.GetFloat("request", "repetition_penalty", 1.0f);
    const float  presence_penalty           = reader.GetFloat("request", "presence_penalty", 0.0f);
    const float  len_penalty                = reader.GetFloat("request", "len_penalty");
    const float  beam_search_diversity_rate = reader.GetFloat("request", "beam_search_diversity_rate");
    const int    min_length                 = reader.GetInteger("request", "min_length", 0);
    const size_t request_batch_size         = reader.GetInteger("request", "request_batch_size");
    // The length of tokens we hope this model to generate
    const int request_output_len = reader.GetInteger("request", "request_output_len");

    FT_CHECK(head_num % tensor_para_size == 0);
    FT_CHECK(decoder_layers % pipeline_para_size == 0);
    FT_CHECK_WITH_INFO(
        repetition_penalty == 1.0f || presence_penalty == 0.0f,
        fmtstr("Found ambiguous parameters repetition_penalty (%f) and presence_penalty (%f) "
               "which are mutually exclusive. Please remove one of repetition_penalty or presence_penalty "
               "or set to a default value.",
               repetition_penalty,
               presence_penalty));

    // Prepare the parallelism parameters
    int rank       = mpi::getCommWorldRank();
    int world_size = mpi::getCommWorldSize();
    if (rank == 0) {
        printf("Total ranks: %d.\n", world_size);
    }
    int device, device_count;
    check_cuda_error(cudaGetDeviceCount(&device_count));
    check_cuda_error(cudaSetDevice(rank % device_count));
    check_cuda_error(cudaGetDevice(&device));

    struct cudaDeviceProp prop;
    check_cuda_error(cudaGetDeviceProperties(&prop, device));
    printf("Device %s\n", prop.name);

    printf("P%d is running with GPU #%d.\n", rank, device);