Add afmoe model
Author: alyosha-swamy
Pull request: yes
Labels: name: New model

## Summary

This PR adds support for the **AFMoE (Arcee Foundational Mixture of Experts)** model architecture for the upcoming **Trinity-Mini** and **Trinity-Nano** releases. AFMoE is a decoder-only transformer model featuring a sparse Mixture of Experts (MoE) approach, combining token-choice routing with shared experts and several architectural innovations for efficient inference and improved performance.

## Model Description

AFMoE features the following key architectural components:

- **Mixture of Experts with Shared Experts**: Combines routed experts (activated per-token via learned routing) with always-active shared experts for stable base computation

- **Token-Choice Routing**: Uses sigmoid or softmax-based routing with normalization and scaling for expert selection

- **Q/K Normalization and Gating**: Applies RMSNorm to query and key projections and uses sigmoid gating on attention outputs for improved training stability

- **Hybrid Attention Patterns**: Alternates between sliding window attention and full attention across layers for efficiency with long contexts

- **Dual Normalization**: Uses pre- and post-normalization around both attention and MLP blocks for training stability

- **Configurable Dense Layers**: Allows initial layers to use dense MLPs before transitioning to sparse MoE layers (`num_dense_layers`)

## Implementation Details

- **Modular implementation** leveraging transformers' modular architecture:

  - Efficient `AfmoeRMSNorm` for layer normalization

  - `AfmoeRotaryEmbedding` for positional encoding

  - `AfmoeAttention` class implementing Q/K normalization and output gating

  - `AfmoeTokenChoiceRouter` for expert selection

  - `AfmoeMoE` class implementing shared + routed experts architecture

  - `AfmoeDecoderLayer` integrating attention and MoE blocks with dual normalization

## Testing

- Added comprehensive test suite following standard transformers test patterns
- Tests for core functionality:
  - Model initialization and weight loading
  - Forward and backward passes
  - Attention mechanism (sliding window + full attention patterns)
  - MoE routing and expert selection
  - RoPE embeddings
  - KV cache compatibility
- Integration tests with example checkpoints
- Verified compatibility with existing transformer infrastructure
- Model loading and inference verified with `arcee-ai/Trinity-Mini`

## Documentation

- Comprehensive model documentation in `docs/source/en/model_doc/afmoe.md`
- Detailed architecture descriptions and usage examples
- All configuration parameters documented with clear descriptions
- Example code for both Pipeline and AutoModel usage patterns

Comments:
- FAILED tests/models/afmoe/test_modeling_afmoe.py::AfmoeModelTest::test_attention_outputs - TypeError: object of type 'NoneType' has no len() 
- FAILED tests/models/afmoe/test_modeling_afmoe.py::AfmoeModelTest::test_prompt_lookup_decoding_matches_greedy_search - TypeError: 'NoneType' object is not subscriptable
- FAILED tests/models/afmoe/test_modeling_afmoe.py::AfmoeModelTest::test_sample_generate_dict_output - AssertionError: Lists differ: [False, False, False] != [True, True, True]
I think the output attention recorder is wrong. Once fixed we can merge

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_42168). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: afmoe, auto

It seems like this model wasn't added to `src/transformers/models/__init__.py` ?

Seems like it - how did the CI pass?