Contribution baseten.co / engine builder: better dataloader in trt-llm

[michaelfeil]

Contribution by baseten.co (running the https://docs.baseten.co/performance/engine-builder-overview engine-builder product)

This PR:

• guarantees OOM-free quantization as fallback, by sequentially lowering the batch size - its a superset of the current "lower by 2 once" low_mem_mode.

In my opinion, this PR should not be required. However the following circumstances have made it necessary:

• https://github.com/NVIDIA/TensorRT-LLM/blob/2ea17cdad28bed0f30e80eea5b1380726a7c6493/tensorrt_llm/quantization/quantize_by_modelopt.py#L458 is currently commented out - the ModelOPT code path is not used.
• max_batch_size is an external / non-default argument.

def quantize_and_export(*,
                        model_dir,
                        device,
                        calib_dataset: str
                        batch_size: int,

at the time of setting this

• to determine the batch_size, we need to load the model. In trt-llm, the model is typically not loaded. Loading the model twice (e.g. 70B llama) would likely lead to a long loading time or an OOM if the model is not freed. A size based on the model safetensors is not possible, as this depends entirely on.g. if the weights will be loaded from fp32 -> fp16

How to solve it:

• tensorrt-llm should use the modelopt workflow - this would not fix the problem, as it does not use batch_size = get_max_batch_size(model, max_sample_length) function yet.
• this PR needs to be implemented in trt-llm or here, best both.

[michaelfeil]

@kevalmorabia97 I see you have reviewed previous PRs here, I'd be glad if you would also review this one.

[kevalmorabia97]

@michaelfeil thanks for your contribution. Let us review this internally and test in our CI. If it looks good we would include this in one of the upcoming releases and give you credit for contribution on the commit

[michaelfeil]

@kevalmorabia97 Any update on this? - no need to credit me, just would love to have this or a similar logic in main.

[cjluo-nv]

Thanks for the contribution @michaelfeil I will merge this to the internal branch and include it in 0.27 release planned for April 3

[michaelfeil]

Thanks!

[kevalmorabia97]

@michaelfeil your contribution is now merged and added you as a co-author of the commit.
Thank you for your contribution. We are working on making more frequent github updates soon.