Adding a mini benchmark for cross entropy loss

[protonu]

This add a benchmark for cross entropy loss where we can vary the vocab size.
There are no linear operators before cross entropy.
The target/labels are pad/sliced to mimic other models we have seen. The input is cast to FP32 and squeeze also to mimic other Hugging Face models we have seen.

[github-actions]

Review updated until commit 0f0c72e

Description

• Added a mini benchmark for cross entropy loss

• Introduced SyntheticMiniModel class for benchmarking

• Expanded vocab size variations for benchmarking

• Included cast and squeeze operations in the model

---

Changes walkthrough 📝

	Relevant files
Enhancement	core.py Improve type hints and IO bytes computation                            benchmarks/python/core.py Added type hint for tensor_props parameter Modified compute_total_iobytes to handle tensor inputs +4/-3      cross_entropy_loss.py Add synthetic model for cross entropy loss benchmark          benchmarks/python/cross_entropy_loss.py Introduced SyntheticMiniModel class Added mini_model, inputs, grads, and generate_vocab_sizes methods +66/-0
Tests	test_cross_entropy_loss.py Add tests for synthetic cross entropy loss benchmark          benchmarks/python/test_cross_entropy_loss.py Imported SyntheticMiniModel Added test_cross_entropy_mini_benchmark_fwd and test_cross_entropy_mini_benchmark_bwd tests +46/-1

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Possible Issue The conditions in lines 158 and 161 seem incorrect. The or condition should likely be and to ensure that inputs and outputs are not torch.Tensor instances. if not isinstance(inputs, Iterable) or isinstance(inputs, torch.Tensor): inputs = [inputs] if not isinstance(outputs, Iterable) or isinstance(outputs, torch.Tensor): outputs = [outputs] Performance Concern The generate_vocab_sizes method generates a large number of vocab sizes, which could lead to a significant increase in benchmarking time. Consider optimizing the number of vocab sizes or adding a mechanism to limit the number of benchmarks. @staticmethod def generate_vocab_sizes(): sizes_from_models = [ 49152, # Starcoder 129280, # DeepSeek-R1 128256, # Llama3 202048, # Llama4 256000, # Gemma2 131072, # Mistral 152064, # Qwen2 32064, # Phi3.5 100352, # Phi4 50264, # GPT-2 ] powers_of_2 = [2**i * 1024 for i in range(4, 9)] combined_set = sorted(set(sizes_from_models) | set(powers_of_2)) # for each vocab size in the set we increment in steps 64 in +/- 5 directions # which gives the total number of vocab sizes to benchmark variations = set() step = 64 for num in combined_set: for i in range(1, 6): variations.add(num + (i * step)) variations.add(num - (i * step)) variations.add(num) return sorted(variations) Hardcoded Batch Size The batch size is hardcoded to 4096 in both test_cross_entropy_mini_benchmark_fwd and test_cross_entropy_mini_benchmark_bwd. Consider making this a parameter or configurable value to allow for more flexible testing. batch_size = 4096 def fwd_call(inp): return SyntheticMiniModel.mini_model(*inp) inputs = SyntheticMiniModel.inputs(int(batch_size), int(vocab_size)) fwd_fn = with_executor(executor, fwd_call) run_benchmark(benchmark, fwd_fn, inputs) @pytest.mark.parametrize( "executor", ["eager", "torchcompile", "thunder", "thunder-torchcompile"] ) @pytest.mark.parametrize("vocab_size", SyntheticMiniModel.generate_vocab_sizes()) def test_cross_entropy_mini_benchmark_bwd(benchmark, executor: str, vocab_size: int): if executor == "torchcompile": clear_dynamo_cache() # picking a value that doesn't OOM for large vocab sizes batch_size = 4096

[jjsjann123]

in line 161 we are wrapping a single tensor output as outputs = [outputs].
So the check here would always be false?

[protonu]

I am not sure if that's quite right.
'''
a = torch.tensor(5)
print(isinstance(a, Iterable))
'''
This should be true.
So outputs could be a single torch tensor and not get wrapped. Please let me know if I am mistaken.

[jjsjann123]

Oh you are right and I didn't know that!

@Priya2698 so it looks like line 158 - 161 isn't doing the right thing?

[jjsjann123]

LGTM

[protonu]

!test

[crcrpar]

Would this dynamic library seriously need to be in this directory?

[protonu]

Thanks for catching that! That was a mistake