diff --git a/common/arg.cpp b/common/arg.cpp index 0d3ddccf..b0904165 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -696,6 +696,13 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, params.next_node_ip = value; } ).set_env("LLAMA_ARG_NEXT_NODE_IP")); + add_opt(llama_arg( + {"--unload", "--unload-weight"}, + format("whether to unload layer weights after use (default: %s)", params.unload ? "true" : "false"), + [](gpt_params & params) { + params.unload = true; + } + ).set_env("LLAMA_ARG_UNLOAD")); add_opt(llama_arg( {"-n", "--predict", "--n-predict"}, "N", format("number of tokens to predict (default: %d, -1 = infinity, -2 = until context filled)", params.n_predict), diff --git a/common/common.cpp b/common/common.cpp index a4a68899..b08f51bc 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1039,6 +1039,7 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param cparams.n_world = params.n_world; cparams.rank = params.rank; cparams.n_layer_window = params.n_layer_window; + cparams.unload = params.unload; if (cparams.master_ip != nullptr) { delete[] cparams.master_ip; diff --git a/common/common.h b/common/common.h index 21540dac..42db03c0 100644 --- a/common/common.h +++ b/common/common.h @@ -147,6 +147,7 @@ struct gpt_params { int32_t n_layer_window = 32; // number of layers to process in each compute std::string master_ip = "localhost"; // ip address of the master node std::string next_node_ip = "localhost"; // ip address of my next node + bool unload = false; // unload layer weights after use or not int32_t n_predict = -1; // new tokens to predict int32_t n_ctx = 0; // context size int32_t n_batch = 2048; // logical batch size for prompt processing (must be >=32 to use BLAS) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index e7678d07..96b8b08c 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2077,17 +2077,19 @@ extern "C" { float wd); // weight decay // graph allocation in a context - GGML_API struct ggml_cgraph * ggml_new_graph (struct ggml_context * ctx); // size = GGML_DEFAULT_GRAPH_SIZE, grads = false - GGML_API struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t size, bool grads); - GGML_API struct ggml_cgraph * ggml_graph_dup (struct ggml_context * ctx, struct ggml_cgraph * cgraph); - GGML_API void ggml_graph_cpy (struct ggml_cgraph * src, struct ggml_cgraph * dst); - GGML_API void ggml_graph_reset (struct ggml_cgraph * cgraph); // set regular grads + optimizer momenta to 0, set loss grad to 1 - GGML_API void ggml_graph_clear (struct ggml_cgraph * cgraph); + GGML_API struct ggml_cgraph * ggml_new_graph (struct ggml_context * ctx); // size = GGML_DEFAULT_GRAPH_SIZE, grads = false + GGML_API struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t size, bool grads); + GGML_API struct ggml_cgraph * ggml_graph_dup (struct ggml_context * ctx, struct ggml_cgraph * cgraph); + GGML_API void ggml_graph_cpy (struct ggml_cgraph * src, struct ggml_cgraph * dst); + GGML_API void ggml_graph_reset (struct ggml_cgraph * cgraph); // set regular grads + optimizer momenta to 0, set loss grad to 1 + GGML_API void ggml_graph_clear (struct ggml_cgraph * cgraph); GGML_API int ggml_graph_size (struct ggml_cgraph * cgraph); GGML_API struct ggml_tensor * ggml_graph_node (struct ggml_cgraph * cgraph, int i); // if i < 0, returns nodes[n_nodes + i] + GGML_API struct ggml_tensor * ggml_graph_leaf (struct ggml_cgraph * cgraph, int i); // if i < 0, returns leafs[n_leafs + i] GGML_API struct ggml_tensor ** ggml_graph_nodes (struct ggml_cgraph * cgraph); GGML_API int ggml_graph_n_nodes(struct ggml_cgraph * cgraph); + GGML_API int ggml_graph_n_leafs(struct ggml_cgraph * cgraph); GGML_API void ggml_graph_add_node(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 03b832d0..6c3720cf 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -19104,6 +19104,16 @@ struct ggml_tensor * ggml_graph_node(struct ggml_cgraph * cgraph, int i) { return cgraph->nodes[i]; } +struct ggml_tensor * ggml_graph_leaf(struct ggml_cgraph * cgraph, int i) { + if (i < 0) { + GGML_ASSERT(cgraph->n_leafs + i >= 0); + return cgraph->leafs[cgraph->n_leafs + i]; + } + + GGML_ASSERT(i < cgraph->n_leafs); + return cgraph->leafs[i]; +} + struct ggml_tensor ** ggml_graph_nodes(struct ggml_cgraph * cgraph) { return cgraph->nodes; } @@ -19112,6 +19122,10 @@ int ggml_graph_n_nodes(struct ggml_cgraph * cgraph) { return cgraph->n_nodes; } +int ggml_graph_n_leafs(struct ggml_cgraph * cgraph) { + return cgraph->n_leafs; +} + void ggml_graph_add_node(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor) { GGML_ASSERT(cgraph->size > cgraph->n_nodes); cgraph->nodes[cgraph->n_nodes] = tensor; diff --git a/include/llama.h b/include/llama.h index b0965123..e1cad355 100644 --- a/include/llama.h +++ b/include/llama.h @@ -318,6 +318,7 @@ extern "C" { uint32_t n_world; // world size uint32_t rank; // my rank uint32_t n_layer_window; // number of layers to process in each compute + bool unload; // whether to unload layer weights after use char * master_ip; // ip address of the master node char * next_node_ip; // ip address of the next node uint32_t n_ctx; // text context, 0 = from model diff --git a/src/llama.cpp b/src/llama.cpp index af62c18f..0b4c12ae 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -1938,7 +1938,7 @@ struct llama_mmap { if (prefetch > 0) { // advise the kernel to preload the mapped memory - if (posix_madvise(addr, std::min(file->size, prefetch), POSIX_MADV_WILLNEED)) { + if (posix_madvise(addr, std::min(file->size, prefetch), POSIX_MADV_WILLNEED)) { LLAMA_LOG_WARN("warning: posix_madvise(.., POSIX_MADV_WILLNEED) failed: %s\n", strerror(errno)); } @@ -1968,6 +1968,10 @@ struct llama_mmap { if (*last <= *first) { *last = *first; } + + GGML_ASSERT(*first % page_size == 0); + GGML_ASSERT(*last % page_size == 0); + GGML_ASSERT(*last >= *first); } // partially unmap the file in the range [first, last) @@ -2562,6 +2566,7 @@ struct llama_cparams { uint32_t n_world; uint32_t rank; uint32_t n_layer_window; + bool unload; uint32_t n_ctx; // context size used during inference uint32_t n_batch; uint32_t n_ubatch; @@ -8935,7 +8940,7 @@ static bool llm_load_tensors( } } - ml.init_mappings(true, use_mlock ? &model.mlock_mmaps : nullptr); + ml.init_mappings(false, use_mlock ? &model.mlock_mmaps : nullptr); // disable prefetch here model.mappings.reserve(ml.mappings.size()); // create the backend buffers @@ -17286,14 +17291,14 @@ static void llama_send_tensors(zmq::socket_t & socket, struct input_tensors * te } static int llama_recv_meta(zmq::socket_t & socket, struct sync_meta * meta) { - socket.setsockopt(ZMQ_RCVTIMEO, 1000); + socket.set(zmq::sockopt::rcvtimeo, 1000); std::vector recv_msgs; if (!zmq::recv_multipart(socket, std::back_inserter(recv_msgs))) { return -1; } - socket.setsockopt(ZMQ_RCVTIMEO, -1); + socket.set(zmq::sockopt::rcvtimeo, -1); for (size_t i = 0; i < recv_msgs.size(); i += 2) { std::string key = recv_msgs[i].to_string(); @@ -17335,6 +17340,106 @@ static void llama_recv_tensors(zmq::socket_t & socket, input_tensors * tensors) } } +static bool is_tensor_loaded(struct ggml_tensor * tensor) { + void * addr = (void *)tensor->data; + size_t size = ggml_nbytes(tensor); + + #ifdef _WIN32 + MEMORY_BASIC_INFORMATION mbi; + size_t bytes_checked = 0; + while (bytes_checked < size) { + if (VirtualQuery((char *)addr + bytes_checked, &mbi, sizeof(mbi)) == 0) { + LLAMA_LOG_ERROR("VirtualQuery failed\n"); + return false; + } + if (mbi.State != MEM_COMMIT) { + return false; // memory not loaded + } + bytes_checked += mbi.RegionSize; + } + return true; + #else + size_t first = (size_t)addr; + size_t last = first + size; + long page_size = sysconf(_SC_PAGESIZE); + + // align addr + llama_mmap::align_range(&first, &last, page_size); + size_t len = std::max(last - first, static_cast(page_size)); + + // calculate the number of pages to check + size_t page_count = (len + page_size - 1) / page_size; + + #ifdef __APPLE__ + char * mincore_res = new char[page_count]; + #else + unsigned char *mincore_res = new unsigned char[page_count]; // use 'unsigned char' for Linux + #endif + + // call mincore to check if pages are resident in memory + if (mincore((void *)first, len, mincore_res) == 0) { + for (size_t i = 0; i < page_count; ++i) { + if (!(mincore_res[i] & 1)) { + delete[] mincore_res; + return false; // page not loaded + } + } + delete[] mincore_res; + return true; // page loaded + } else { + LLAMA_LOG_ERROR("mincore failed\n"); + delete[] mincore_res; + return false; + } + #endif +} + +static float is_graph_loaded(struct ggml_cgraph * cgraph) { + uint32_t n_loaded = 0; + uint32_t n_total = 0; + for (int i = 0; i < ggml_graph_n_leafs(cgraph); i++) { + struct ggml_tensor * cur = ggml_graph_leaf(cgraph, i); + if (strstr(cur->name, "weight") == nullptr || cur->data == nullptr) { + continue; + } + if (is_tensor_loaded(cur)) n_loaded++; + n_total++; + } + return float(n_loaded) / float(n_total) * 100.0f; +} + +static void manage_graph_tensors(struct ggml_cgraph * cgraph, int advice, bool force = false) { + for (int i = 0; i < ggml_graph_n_leafs(cgraph); i++) { + struct ggml_tensor * cur = ggml_graph_leaf(cgraph, i); + + if (strstr(cur->name, "weight") == nullptr || cur->data == nullptr) { + continue; + } + + void * addr = (void *)cur->data; + size_t size = ggml_nbytes(cur); + size_t first = (size_t)addr; + size_t last = first + size; + long page_size = sysconf(_SC_PAGESIZE); + + // align addr + llama_mmap::align_range(&first, &last, page_size); + size_t len = std::max(last - first, static_cast(page_size)); + + // hint to load memory + posix_madvise((void *)first, len, advice); + + // if advice is POSIX_MADV_WILLNEED, force to prefetch data + if (force && advice == POSIX_MADV_WILLNEED) { + volatile char * ptr = (volatile char *)first; + for (size_t off = 0; off < len; off += page_size) { + volatile char data = ptr[off]; + (void)data; + } + } + } +} + // decode a batch of tokens by evaluating the transformer // // - lctx: llama context @@ -17533,6 +17638,8 @@ static int llama_decode_internal( ggml_backend_sched_alloc_graph(lctx.sched[i], gf[i]); } + ggml_cgraph * sub_gf = nullptr; + ggml_cgraph * next_gf = nullptr; const uint32_t n_layer = hparams.n_layer; const char * layer_str = nullptr; int cur_l = -1; @@ -17541,7 +17648,8 @@ static int llama_decode_internal( GGML_ASSERT(my_rank == 0 || n_world > 1); for (size_t i = 0; i < (size_t)gf.size(); ++i) { - ggml_cgraph * sub_gf = gf[i]; + sub_gf = gf[i]; + next_gf = gf[(i + 1) % gf.size()]; if (n_world > 1 && !(my_rank == 0 && i == 0) && !(my_rank == 0 && is_last_l)) { { // receive data from previous nodes @@ -17600,6 +17708,17 @@ static int llama_decode_internal( zmq::socket_t * s = is_to_master ? lctx.master_socket : lctx.send_socket; llama_send_tensors(*s, &tensors); } + + // overlap memory scheduling with other nodes' communication and computing + if (cparams.unload) { + timer(manage_graph_tensors); + if (n_world != 1) { + manage_graph_tensors(sub_gf, POSIX_MADV_DONTNEED); + if (!(my_rank == 0 && is_last_l)) { + manage_graph_tensors(next_gf, POSIX_MADV_WILLNEED, true); + } + } + } } // update the kv ring buffer @@ -19314,6 +19433,7 @@ struct llama_context_params llama_context_default_params() { /*.n_world =*/ 1, /*.rank =*/ 0, /*.n_layer_window =*/ 32, + /*.unload =*/ false, /*.master_ip =*/ nullptr, /*.next_node_ip =*/ nullptr, /*.n_ctx =*/ 512, @@ -19617,6 +19737,7 @@ struct llama_context * llama_new_context_with_model( cparams.n_world = params.n_world; cparams.rank = params.rank; cparams.n_layer_window = params.n_layer_window; + cparams.unload = params.unload; cparams.n_seq_max = std::max(1u, params.n_seq_max); cparams.n_threads = params.n_threads; cparams.n_threads_batch = params.n_threads_batch; @@ -19995,6 +20116,9 @@ struct llama_context * llama_new_context_with_model( std::vector gf = llama_build_graph(*ctx, ubatch, true); ctx->sched.resize(gf.size()); + + // prefetch the first subgraph weights + manage_graph_tensors(gf.front(), POSIX_MADV_WILLNEED, true); // initialize scheduler with the worst-case graph bool ok = true;