Upload FlaxLlamaForCausalLM

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+easydel-model.parameters filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,79 @@
+---
+tags:
+- EasyDeL
+- FlaxLlamaForCausalLM
+- safetensors
+- TPU
+- GPU
+- XLA
+- Flax
+---
+# EasyDeL/EasyDeL-Llama-3.2-1B-Instruct
+
+[![EasyDeL](https://img.shields.io/badge/🤗_EasyDeL-0.0.80-blue.svg)](https://github.com/erfanzar/EasyDeL)
+[![Model Type](https://img.shields.io/badge/Model_Type-FlaxLlamaForCausalLM-green.svg)](https://github.com/erfanzar/EasyDeL)
+
+A model implemented using the EasyDeL framework, designed to deliver optimal performance for large-scale natural language processing tasks.
+
+## Overview
+
+EasyDeL provides an efficient, highly-optimized, and customizable machine learning model compatible with both GPU and TPU environments. Built with JAX, this model supports advanced features such as sharded model parallelism, making it suitable for distributed training and inference and customized kernels.
+
+## Features
+
+
+- **Efficient Implementation**: Built with JAX/Flax for high-performance computation.
+- **Multi-Device Support**: Optimized to run on TPU, GPU, and CPU environments for sharding model over 2^(1-1000+) of devices.
+- **Sharded Model Parallelism**: Supports model parallelism across multiple devices for scalability.
+- **Customizable Precision**: Allows specification of floating-point precision for performance optimization.
+
+
+## Installation
+
+To install EasyDeL, simply run:
+
+```bash
+pip install easydel
+```
+
+## Usage
+
+### Loading the Pre-trained Model
+
+To load a pre-trained version of the model with EasyDeL:
+
+```python
+from easydel import AutoEasyDeLModelForCausalLM
+from jax import numpy as jnp, lax
+
+max_length = None # can be set to use lower memory for caching
+
+# Load model and parameters
+model, params = AutoEasyDeLModelForCausalLM.from_pretrained(
+    "EasyDeL/EasyDeL-Llama-3.2-1B-Instruct",
+    config_kwargs=dict(
+        use_scan_mlp=False,
+        attn_dtype=jnp.float16,
+        freq_max_position_embeddings=max_length,
+        mask_max_position_embeddings=max_length,
+        attn_mechanism=ed.AttentionMechanisms.FLASH_ATTN2
+    ), 
+    dtype=jnp.float16,
+    param_dtype=jnp.float16,
+    precision=lax.Precision("fastest"),
+    auto_shard_params=True, 
+)
+```
+
+## Supported Tasks
+
+
+[Need more information]
+
+ 
+## Limitations
+
+
+- **Hardware Dependency**: Performance can vary significantly based on the hardware used.
+- **JAX/Flax Setup Required**: The environment must support JAX/Flax for optimal use.
+- **Experimental Features**: Some features (like custom kernel usage or ed-ops) may require additional configuration and tuning.

config.json

@@ -0,0 +1,100 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_axis_name": "sp",
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attn_mechanism": "vanilla",
+  "axis_dims": [
+    1,
+    1,
+    1,
+    -1
+  ],
+  "axis_names": [
+    "dp",
+    "fsdp",
+    "tp",
+    "sp"
+  ],
+  "backend": null,
+  "bits": null,
+  "block_b": 1,
+  "block_k": 128,
+  "block_q": 128,
+  "bos_token_id": 128000,
+  "easy_method": "train",
+  "embd_pdrop": 0.0,
+  "eos_token_id": [
+    128001,
+    128008,
+    128009
+  ],
+  "fcm_max_ratio": 0.0,
+  "fcm_min_ratio": 0.0,
+  "flash_attention_backward_pass_impl": "triton",
+  "freq_max_position_embeddings": 32768,
+  "gradient_checkpointing": "nothing_saveable",
+  "hardware_abstraction": false,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "mask_max_position_embeddings": 32768,
+  "max_position_embeddings": 131072,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 8,
+  "number_rep_kv": 1,
+  "pallas_k_block_size": null,
+  "pallas_m_block_size": null,
+  "pallas_n_block_size": null,
+  "partition_axis": [
+    [
+      "fsdp",
+      "dp"
+    ],
+    "sp",
+    "sp",
+    "tp",
+    "sp",
+    "tp",
+    null,
+    null,
+    null,
+    null,
+    "tp",
+    "sp",
+    null
+  ],
+  "platform": null,
+  "pretraining_tp": 1,
+  "quantize_kv_cache": false,
+  "resid_pdrop": 0.0,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": {
+    "factor": 32.0,
+    "high_freq_factor": 4.0,
+    "low_freq_factor": 1.0,
+    "original_max_position_embeddings": 8192,
+    "rope_type": "llama3"
+  },
+  "rope_theta": 10000.0,
+  "scan_attention_layers": false,
+  "scan_layers": true,
+  "scan_mlp_chunk_size": 1024,
+  "scan_ring_attention": true,
+  "shard_attention_computation": true,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.46.2",
+  "use_cache": true,
+  "use_scan_mlp": false,
+  "use_sharded_kv_caching": true,
+  "use_sharding_constraint": false,
+  "vocab_size": 128256
+}

easydel-model.parameters

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2467e3fb9a96c3736fdb4940b934135634c14eea0b89ddb54df6de73fca50420
+size 5993947552

generation_config.json

@@ -0,0 +1,12 @@
+{
+  "bos_token_id": 128000,
+  "do_sample": true,
+  "eos_token_id": [
+    128001,
+    128008,
+    128009
+  ],
+  "temperature": 0.6,
+  "top_p": 0.9,
+  "transformers_version": "4.46.2"
+}