Test different input sequence lengths for Llama #1070
Conversation
pmarkovicTT
requested review from
nvukobratTT,
pilkicTT and
dgolubovicTT
as code owners
| ], | ||
| ) | ||
| @pytest.mark.parametrize("seq_len", [1, 2, 4, 7, 8, 16, 28, 32, 63, 64, 99, 117, 128, 256, 341, 512, 1024, 1790, 2048]) | ||
| @pytest.mark.skip(reason="No need to run in CI as it takes a long time to run.") |
There was a problem hiding this comment.
My recommendation is to choose which of these will be part of the training focus, instead of skipping it entirely.
E.g. if we're going to focus on training 2048 seq len model, let's fully compile and run as part of push CI that variant alone.
There was a problem hiding this comment.
That's right - understanding which sequence length is relevant for Llama finetuning is one of the training team's tasks.
Once we establish which set of seq lengths is needed, we will continue with PCC tests and run as part of CI.
There was a problem hiding this comment.
Agree as well, will update seq_len parameters with required ones for training once we choose them (we will run some experiments separately).
There was a problem hiding this comment.
This is updated to use only dim sizes we care about. Additionally, I setup only one hidden layer to be used for test to speed it up (while I also ran full model test locally to make sure it passes).
| input_ids = tokenizer(prompt, padding="max_length", truncation=True, return_tensors="pt").input_ids | ||
|
|
||
| # Compile the model and run fwd pass | ||
| compiled_model = forge.compile(framework_model, input_ids) |
There was a problem hiding this comment.
Do we want to test out bwd compile/run as well?
One general question, is there a clean way to test a backward part of a graph in isolation? For example, our compile should return compiled context that contains information about each compiled component (e.g. fwd, bwd, loss, etc.).
Therefore, is there a clean way to just call the bwd part of the graph with random inputs, without a need to run the forward part, and initialize the loss and optimizer part of the training workflow?
Note: this is not a requirement for this PR, just a general question that can be useful here as well. I.e. can we have granular tests that target specific functionality, rather than the whole workflow (only the bwd part of the model). I see this as especially useful for bwd generallity push in the future. cc @vladimirjovanovicTT
There was a problem hiding this comment.
I think this is a must-have functionality as part of our training generality/BFS effort.
Let's discuss the implementation details offline.
|
||||||||||||||
|
||||||||||||||
|
||||||||||||||
|
||||||||||||||
pmarkovicTT
force-pushed
the
pmarkovic/test-input-seq-lengths-llama
branch
from
5b45490 to
710afb4
Compare
|
||||||||||||||
1 similar comment
|
||||||||||||||
|
||||||||||||||
1 similar comment
|
||||||||||||||
| "model_path", | ||
| [ | ||
| "openlm-research/open_llama_3b", | ||
| pytest.param("meta-llama/Llama-3.2-1B", marks=pytest.mark.xfail(reason="Unsupported Op: repeat_interleave")), |
There was a problem hiding this comment.
Support for repeat_interleave is added, so feel free to test out 3.2. 1B on latest main :))
|
|
||
| prompt = "Q: What is the largest animal?\nA:" | ||
| input_ids = tokenizer(prompt, padding="max_length", truncation=True, return_tensors="pt").input_ids | ||
| input_ids = input_ids.to(torch.int32) |
There was a problem hiding this comment.
Why is this required? What is the default type for input IDs?
Do we expect that embedding input will always be int-based? If yes, maybe we should have a pass that will encompass this.
There was a problem hiding this comment.
Default type is int64 and we need to cast it due to following issue #952
Yep, embedding inputs are int-based (indices in the vocabulary), but I am not sure what you mean about another pass.
| @@ -16,6 +16,7 @@ def load_model(model_path="openlm-research/open_llama_3b", **kwargs): | |||
| config.use_cache = kwargs.get("use_cache", False) | |||
| config.output_attentions = kwargs.get("output_attentions", False) | |||
| config.output_hidden_states = kwargs.get("output_hidden_states", False) | |||
| config.num_hidden_layers = kwargs.get("num_hidden_layers", 26) | |||
There was a problem hiding this comment.
Was this intentional?
Any specific reasons for updating original number of hidden layers?
There was a problem hiding this comment.
Yep, that's per our discussion in the last sync. Running llama with all layers takes quite some time and since this is not e2e/demo test, I thought it makes sense to speed it up by using a single layer.
| pytest.param("meta-llama/Llama-3.2-1B", marks=pytest.mark.xfail(reason="Unsupported Op: repeat_interleave")), | ||
| ], | ||
| ) | ||
| @pytest.mark.parametrize("seq_len", [128, 512, 2048]) |
There was a problem hiding this comment.
Any thoughts on testing on lower precisions? E.g. bfloat16?
In full precision, Open Llama will require 12GB, while 3.2. 4GB. That said, we should either:
- Test out lower precision DF (Open Llama will barely fit n150 for inference, definitely not for training)
- Focus only on Llama 3.2 for training. In this case as well, we'll need to run in half-precision for training in order to fit on n150 (depending on which optimizer we use during fine-tuning; full training will probably be a stretch)
There was a problem hiding this comment.
Yep you are completely right. That's something we plan to do with llama backward/training tests, and eventually incorporate this test there as well. We are currently investigating memory footprint of llama models on GPU to find optimal setup for our devices. Our plan is to add tests based on findings.
|
@pmarkovicTT is this one ready for review and potential merge? If not, can we move it to draft? |
pmarkovicTT
force-pushed
the
pmarkovic/test-input-seq-lengths-llama
branch
from
710afb4 to
f081375
Compare
|
||||||||||||||
|
||||||||||||||
|
||||||||||||||||||
|
||||||||||||||
|
@nvukobratTT PR is ready for review/merge. Summary of previous conversations:
|
|
||||||||||||||
|
||||||||||||||
|
||||||||||||||
|
||||||||||||||
Add test to make sure Llama compiles and run fwd pass with different input sequence lengths as we will have inputs of various lengths during training.
Close #1071