Introduction to SageMaker JumpStart - Text Generation with Falcon models

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

In this demo notebook, we demonstrate how to use the SageMaker Python SDK to deploy Falcon models for text generation. It is a permissively licensed (Apache-2.0) open source model trained on the RefinedWeb dataset. We show several example use cases including code generation, question answering, translation etc.

Changing instance type

Models have been tested on the following instance types:

• Falcon 7B and 7B instruct: ml.g5.2xlarge, ml.g5.2xlarge, ml.g5.4xlarge, ml.g5.8xlarge, ml.g5.16xlarge, ml.g5.12xlarge, ml.g5.24xlarge, ml.g5.48xlarge, ml.p4d.24xlarge
• Falcon 40B and 40B instruct: ml.g5.12xlarge, ml.g5.48xlarge, ml.p4d.24xlarge

If an instance type is not available in you region, please try a different instance. You can do so by specifying instance type in the JumpStartModel class.

my_model = JumpStartModel(model_id="huggingface-llm-falcon-40b-instruct-bf16", instance_type="ml.g5.12xlarge")

Changing number of GPUs

Falcon models are served with HuggingFace (HF) LLM DLC which requires specifying number of GPUs during model deployment.

Falcon 7B and 7B instruct: HF LLM DLC currently does not support sharding for 7B model. Thus, even if more than one GPU is available on the instance, please do not increase number of GPUs.

Falcon 40B and 40B instruct: By default number of GPUs are set to 4. However, if you are using ml.g5.48xlarge or ml.p4d.24xlarge, you can increase number of GPUs to be 8 as follows:

my_model = JumpStartModel(model_id="huggingface-llm-falcon-40b-instruct-bf16", instance_type="ml.g5.48xlarge")

my_model.env['SM_NUM_GPUS'] = '8'

predictor = my_model.deploy()

About the model

Falcon is a causal decoder-only model built by Technology Innovation Institute (TII) and trained on more than 1 trillion tokens of RefinedWeb enhanced with curated corpora. It was built using custom-built tooling for data pre-processing and model training built on Amazon SageMaker. As of June 6, 2023, it is the best open-source model currently available. Falcon-40B outperforms LLaMA, StableLM, RedPajama, MPT, etc. To see comparison, see OpenLLM Leaderboard. It features an architecture optimized for inference, with FlashAttention and multiquery.

Refined Web Dataset: Falcon RefinedWeb is a massive English web dataset built by TII and released under an Apache 2.0 license. It is a highly filtered dataset with large scale de-duplication of CommonCrawl. It is observed that models trained on RefinedWeb achieve performance equal to or better than performance achieved by training model on curated datasets, while only relying on web data.

Model Sizes:

• Falcon-7b: It is a 7 billion parameter model trained on 1.5 trillion tokens. It outperforms comparable open-source models (e.g., MPT-7B, StableLM, RedPajama etc.). To see comparison, see OpenLLM Leaderboard. To use this model, please select model_id in the cell above to be "huggingface-textgeneration-falcon-7b-bf16".
• Falcon-40B: It is a 40 billion parameter model trained on 1 trillion tokens. It has surpassed renowned models like LLaMA-65B, StableLM, RedPajama and MPT on the public leaderboard maintained by Hugging Face, demonstrating its exceptional performance without specialized fine-tuning. To see comparison, see OpenLLM Leaderboard.

Instruct models (Falcon-7b-instruct/Falcon-40B-instruct): Instruct models are base falcon models fine-tuned on a mixture of chat and instruction datasets. They are ready-to-use chat/instruct models. To use these models, please select model_id in the cell above to be "huggingface-textgeneration-falcon-7b-instruct-bf16" or "huggingface-textgeneration-falcon-40b-instruct-bf16".

It is recommended that Instruct models should be used without fine-tuning and base models should be fine-tuned further on the specific task.

Limitations:

• Falcon models are mostly trained on English data and may not generalize to other languages.
• Falcon carries the stereotypes and biases commonly encountered online and in the training data. Hence, it is recommended to develop guardrails and to take appropriate precautions for any production use. This is a raw, pretrained model, which should be further finetuned for most usecases.

Supported parameters

Some of the supported parameters while performing inference are the following:

• max_length: Model generates text until the output length (which includes the input context length) reaches max_length. If specified, it must be a positive integer.
• max_new_tokens: Model generates text until the output length (excluding the input context length) reaches max_new_tokens. If specified, it must be a positive integer.
• num_beams: Number of beams used in the greedy search. If specified, it must be integer greater than or equal to num_return_sequences.
• no_repeat_ngram_size: Model ensures that a sequence of words of no_repeat_ngram_size is not repeated in the output sequence. If specified, it must be a positive integer greater than 1.
• temperature: Controls the randomness in the output. Higher temperature results in output sequence with low-probability words and lower temperature results in output sequence with high-probability words. If temperature -> 0, it results in greedy decoding. If specified, it must be a positive float.
• early_stopping: If True, text generation is finished when all beam hypotheses reach the end of sentence token. If specified, it must be boolean.
• do_sample: If True, sample the next word as per the likelihood. If specified, it must be boolean.
• top_k: In each step of text generation, sample from only the top_k most likely words. If specified, it must be a positive integer.
• top_p: In each step of text generation, sample from the smallest possible set of words with cumulative probability top_p. If specified, it must be a float between 0 and 1.
• return_full_text: If True, input text will be part of the output generated text. If specified, it must be boolean. The default value for it is False.
• stop: If specified, it must a list of strings. Text generation stops if any one of the specified strings is generated.

We may specify any subset of the parameters mentioned above while invoking an endpoint.

For more parameters and information on HF LLM DLC, please see this article.

Limits on the number of input and output tokens

Large models such as Falcon have very high accelerator memory footprint. Thus, a very large input payload or generating a large output can cause out of memory errors. Furthermore, generating large outputs can take secs or even minutes. However, SageMaker has a response time limit of 60 seconds. Thus, large input or output payload can cause timeout issues. Based on these two constraints, we recommend the following limits on the input and new tokens

Note that, limits don't apply equally to input tokens and new tokens. Models typically support much larger input tokens if you decrease max_new_tokens slightly. Also, note that non-default values of inference parameters will impact the size of input and output payload supported. For instance, higher value of num_beams will reduce the number of max_new_tokens you can generate.

Words-Token ratio: Ratio of words to tokens is roughly 1.5. So, if number of input tokens 900, it corresponds to ~600 input words. Note that this is not always the case. There are several pieces of text where words-tokens ratio can be significantly different.

Setting non-default environment variables: If setting num_input_tokens >1024 or num_total_tokens >=2048, you would need to change the environment variable before deploying the model: my_model.env['MAX_INPUT_LENGTH'] = '2048' (default '1024') my_model.env['MAX_TOTAL_TOKENS'] = '4096' (default '2048')

Note that the endpoint supports maximum length on the number of total tokens (number of input + new tokens) whereas we report the limits above on number of input tokens and number of new tokens.

If using ml.g5.48xlarge for Falcon 40b, you would need to use 8 GPUs (my_model.env['SM_NUM_GPUS'] = '8')

Concurrency/invoking at short intervals: When sending multiple requests at once or at very short interval, the endpoint can not handle the same thresholds as above. It may result in OOM CUDA error. Thus, we recommend trying smaller number of input and output tokens than mentioned above. This is because the TGI container batches requests. While this results in significant throughput improvements, it may cause CUDA OOM or timeout errors. To reduce the number of batched requests, you can set the max_concurrent_requests parameter (my_model.env['MAX_CONCURRENT_REQUESTS'] = '1').

Corrupted endpoint state: It has been observed that once an endpoint suffers a CUDA OOM error, the endpoint may get in corrupted state where it may not function for very small inputs and small max_new_token parameter. It has been observed that letting endpoint sit ideally can help reset the state. If that does not work, please restart the endpoint (delete and launch a new endpoint).

Model quality: Note that above limits only refer to the memory limitation of the instance types available and the sagemaker endpoint response timeout limit. Model itself can theoretically support arbitrary large input and output payload if infinite CUDA memory is available and there is no limit on sagemaker endpoint response time. However, it has been observed that the quality of the model's output decrease substantially when provided a very large input payload (eg. generate summary of a document with 20,000 tokens) or generate large outputs (write a story with 2000 tokens). Thus, you may want to stay below the limits recommended above to generate high quality outputs.

Generating few tokens at a time - Supporting large outputs on smaller instances

As observed above, model can support much larger input tokens than output tokens. For instance, endpoint has higher latency and more memory requirement when given an input sequence of length 100 tokens and generating 100 tokens compared to the case with input sequence of length 190 tokens and generating 10 new tokens. Based on this observation, we can set up our text generation process to avoid CUDA OOM issue and endpoint response timeout issue by invoke endpoint repeatedly to generate very large output sequence which would have been otherwise infeasible. For Falcon 40b instruct model, we observed that you can generate more than 5000 new tokens on ml.g5.12xlarge if you generate 100 tokens at a time.

We also experimented with the computation overhead of repeatedly querying endpoint and thus computing activations/states for the input text mutliple times. We observed that even when generating 10 tokens at a time, this contributed less than 5% of the overall time to generate the desired number of output tokens. Thus, even when generating the entire sequence is feasible for a single query, you can simply generate generate in batches.

5. Clean up the endpoint