llama3 rope

[RaymondLi0]

✨ Description

Closes #39

🔍 Type of change

Select all that apply:

• [ ] 🐛 Bug fix (non-breaking change that addresses a specific issue)
• [x] 🚀 New feature (non-breaking change that adds functionality)
• [ ] ⚠️ Breaking change (a change that could affect existing functionality)
• [ ] 📈 Performance improvement/optimization (improves speed, memory usage, or efficiency)
• [ ] 🛠️ Code refactor (non-functional changes that improve code readability, structure, etc.)
• [ ] 📦 Dependency bump (updates dependencies, including Dockerfile or package changes)
• [ ] 📝 Documentation change (updates documentation, including new content or typo fixes)
• [ ] 🔧 Infrastructure/Build change (affects build process, CI/CD, or dependencies)

📝 Changes

• add llama3-style rope scaling

✅ Checklist

Make sure the following tasks are completed before submitting the PR:

General

• [ ] 📜 I have read and followed the contributing guidelines.
• [ ] 🏷️ I am using a clear and descriptive PR title that summarizes the key change or feature introduced.
• [ ] 🎉 The functionality is complete, and I have tested the changes.
• [ ] 📝 I have updated the documentation if needed.
• [ ] ⚠️ The change does not introduce any new issues (e.g., runtime warnings, type checker errors, linting problems, unhandled edge cases).
• [ ] 🧩 I have commented my code, especially in hard-to-understand areas.

Dependencies and Configuration

• [ ] 🐋 I have updated the Docker configuration or dependencies, if applicable.
• [ ] 🔄 I have ensured compatibility with the existing setup after dependency changes.

Testing

• [ ] 🧪 I have added or updated tests to cover my changes.
• [ ] ✔️ New and existing tests pass locally with my changes.
• [ ] 🚦 I have tested these changes on GPUs and verified training stability.
• [ ] 🏋️ I have tested the changes on realistic training workloads, if applicable.

Performance Impact

• [ ] 📊 I have run benchmarks where applicable to evaluate the performance impact.
• [ ] ✅ The benchmarks show no performance regression.
• [ ] 🚀 The benchmarks indicate a potential performance improvement.
• [ ] ⚠️ The benchmarks indicate a potential performance degradation.
• [ ] 📈 I have provided benchmark results and detailed any performance impact below, if applicable.

📊 Performance Impact Details

If there is any impact on performance, describe it and provide benchmark results, if applicable:

---

🗒️ Additional Notes

TODOs:

• [ ] Implement HF conversion

[tscholak]

Hi @RaymondLi0! Functionally this looks like what we want (pending model conversion), but are you confident (i.e. have you checked) that the forward and backward passes of hf-llama and fast-llm-llama are the same?

[RaymondLi0]

Hi @RaymondLi0! Functionally this looks like what we want (pending model conversion), but are you confident (i.e. have you checked) that the forward and backward passes of hf-llama and fast-llm-llama are the same?

Haven't done that check. Do we have existing tests comparing forward/backward of fast-llm and hf-transformers? If no I can look into adding this

[jlamypoirier]

Haven't done that check. Do we have existing tests comparing forward/backward of fast-llm and hf-transformers? If no I can look into adding this

There is one in test_checkpoint https://github.com/ServiceNow/Fast-LLM/blob/main/tests/test_checkpoint.py#L31. It could work for this one if we added a llama3 model to the testing suite (in common.py). (It would also be a good test of conversion, etc.)