The model generate repeated words.

Describe the bug
I use the example code provided on the documentation (https://erfanzar.github.io/EasyDeL/Llama2/).
But I changed the model to 'hfl/chinese-alpaca-2-7b-16k'. (The model is based on llama-2-chat and has a linear rope scale in config.json)
Then when I run the code. the model repeated to generate
“Sharding Params: 100%|████████████████████████| 291/291 [02:03<00:00,  2.36it/s]
Compiling NonGreedy(Generate) Funcs
. Beijing. Beijing in English.ing in Chinese.ing Beijing Beijing Beijing Beijing Beijing Beijing Beijing Beijing Beijing Beijing Beijing, China Daily Beijing Beijing. Beijing Beijinginginginginginginginginginginginginginginginging Beijinginging Beijing Beijing Beijing Beijing Beijing Beijinging Beijingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingingi”
How can I solve above issue? Thank you so much. this is really an amazing work.
To Reproduce
Below is my Full code.
`from EasyDel.modules import FlaxLlamaForCausalLM
from EasyDel.serve import JAXServer
import jax
from EasyDel.transform import llama_from_pretrained
from transformers import AutoTokenizer
model_id = 'hfl/chinese-alpaca-2-7b-16k'
params, config = llama_from_pretrained(
model_id,
device=jax.devices('cpu')[0]  # Offload on CPU
)
server = JAXServer.load_from_params(
params=params,
model=FlaxLlamaForCausalLM(
config=config,
dtype=jax.numpy.bfloat16,  # Im on TPUs
param_dtype=jax.numpy.bfloat16,  # Im on TPUs
precision=jax.lax.Precision('fastest'),
_do_init=False,
input_shape=(1, 1024)
),
config_model=config,
add_params_field=True,
tokenizer=AutoTokenizer.from_pretrained(model_id),
verbose=False,
do_memory_log=True,
)
response_printed = 0
for response, tokens_used in server.process(
'Please introduce Beijing in detail', stream=True
):
print(response[response_printed:], end='')
response_printed = len(response)`

If the model has linear or custom rope that might be the cause ill fix that in 10 hours
But in case that you are not using any special thing like that can i know number of your devices and your mesh size?

thanks for your reply. I ran above code on the Kaggle using TPU VM v3-8.
in addition, i didn't find any parameter about repeat_penalty. if i want to change this parameter to a specific value like 1.15. how can i achieve that?
thanks!

it's probebly fixed by now you can run your code again

@erfanzar I tried just now. the issue still exists. e.g.
Sharding Params: 100%|████████████████████████| 291/291 [02:18<00:00,  2.11it/s]
Compiling NonGreedy(Generate) Funcs
######################
. Points.ing.ing is there are there is there are you can you should you are there are there is Beijing is there is there is there is there is there is there is there is there is there are there is there are there are there is there is there is there is there is there is there are there is there is there is there is there is there are there is there are there are there is there is there is there is there is there is there are there is there is there is there is there is there are there is there are there are there are there is there is there is there is there is there are there is there is there is there are there is there are there is there are there are there are there is there is there is there is there is there are there are there are there is there are there is there are there is there are there are there are there is there is there is there is there are there are there are there are there is there are there is there are there are there are there are there are there is there is there is there are there are there are there are there are there are.
`%cd /root
!apt-get update && apt-get upgrade -y
!apt-get install golang -y
!pip uninstall tensorflow -y
!pip install -U transformers pymongo transformers sentencepiece accelerate
!git clone https://github.com/erfanzar/EasyDeL.git
!cd EasyDeL && pip install -r requirements.txt
!pip install git+https://github.com/erfanzar/EasyDeL.git
from EasyDel.modules import FlaxLlamaForCausalLM
from EasyDel.serve import JAXServer,JAXServerConfig
import jax
from EasyDel.transform import llama_from_pretrained
from transformers import AutoTokenizer
model_id = 'hfl/chinese-alpaca-2-7b-16k'
params, config = llama_from_pretrained(
model_id,
device=jax.devices('cpu')[0]  # Offload on CPU
)
configs = JAXServerConfig(
contains_auto_format=False,
max_length=1024,
max_new_tokens=1024,
max_stream_tokens=32,
temperature=0.8,
top_p=0.9,
top_k=40,
logging=False,
use_prefix_tokenizer=True
)
server = JAXServer.load_from_params(
params=params,
model=FlaxLlamaForCausalLM(
config=config,
dtype=jax.numpy.bfloat16,  # Im on TPUs
param_dtype=jax.numpy.bfloat16,  # Im on TPUs
precision=jax.lax.Precision('fastest'),
_do_init=False,
input_shape=(1, 1024)
),
config_model=config,
add_params_field=True,
tokenizer=AutoTokenizer.from_pretrained(model_id),
verbose=False,
do_memory_log=True,
config=configs
)
response_printed = 0
for response, tokens_used in server.process(
'Please introduce Beijing in detail', stream=True
):
print(response[response_printed:], end='')
response_printed = len(response)`
Above is my full code on kaggle tpu. What am I missing? Thank you.

I used the same code and tried another models and the outputs were not normal either.

chinese-alpaca-2-7b and chinese-alpaca-2-7b-16k. (the only difference between the two are rope scaling.)
togethercomputer/Llama-2-7B-32K-Instruct

it seems like there are some other bugs in the code?

the model and base have no problem I have tested the model and checked the model with the llama_flax.py test and it performed completely the same result

it must be due to wrong partitioning since the model you are using has 50K~ vocab size and maybe that's not being shared correctly

if you are not an expert in writing partition Specs in Jax use this one
partition_rules = (
            ("model/embed_tokens/embedding", PS("tp", ("fsdp", "sp"))),
            ("self_attn/(q_proj|k_proj|v_proj)/kernel", PS(("fsdp", "sp"), "tp")),
            ("self_attn/o_proj/kernel", PS("tp", ("fsdp", "sp"))),

            ("mlp/gate_proj/kernel", PS(("fsdp", "sp"), "tp")),
            ("mlp/down_proj/kernel", PS("tp", ("fsdp", "sp"))),
            ("mlp/up_proj/kernel", PS(("fsdp", "sp"), "tp")),

            ("input_layernorm/kernel", PS(None)),
            ("post_attention_layernorm/kernel", PS(None)),

            ("model/norm/kernel", PS(None)),
            ("lm_head/kernel", PS(("fsdp", "sp"), "tp")),
            ('.*', PS(None))
)

I tried it with the original code and llama-2-7b-chat and it is also repeating itself.
will provide a minimal example soon.

ok, it was much simpler. just try to generate with this:
python -m examples.serving.causal-lm.llama-2-chat   --repo_id='meta-llama/Llama-2-7b-chat-hf' --max_length=4096   --max_new_tokens=2048 --max_stream_tokens=32 --temperature=0.6   --top_p=0.95 --top_k=50   --dtype='fp16' --use_prefix_tokenizer

@erfanzar sorry, I'm not an expert in programing. I might do it wrong. I don't know where to put your partition_rules.
so I tried below code. but it return error when generating.
`server = JAXServer.load_from_params(
params=params,
model=model,
config_model=config,
add_params_field=True,
tokenizer=tokenizer,
verbose=False,
do_memory_log=True,
config=configs
)
partition_rules = (
("model/embed_tokens/embedding", PS("tp", ("fsdp", "sp"))),
("self_attn/(q_proj|k_proj|v_proj)/kernel", PS(("fsdp", "sp"), "tp")),
("self_attn/o_proj/kernel", PS("tp", ("fsdp", "sp"))),
        ("mlp/gate_proj/kernel", PS(("fsdp", "sp"), "tp")),
        ("mlp/down_proj/kernel", PS("tp", ("fsdp", "sp"))),
        ("mlp/up_proj/kernel", PS(("fsdp", "sp"), "tp")),

        ("input_layernorm/kernel", PS(None)),
        ("post_attention_layernorm/kernel", PS(None)),

        ("model/norm/kernel", PS(None)),
        ("lm_head/kernel", PS(("fsdp", "sp"), "tp")),
        ('.*', PS(None))

)
server.auto_configure(model,params,tokenizer,partition_rules)
response_printed = 0
for response, tokens_used in server.process(
'Please introduce Beijing in detail', stream=True
):
print(response[response_printed:], end='')
response_printed = len(response)
`

ScopeCollectionNotFound                   Traceback (most recent call last)
Cell In[4], line 3
1 print('######################')
2 response_printed = 0
----> 3 for response, tokens_used in server.process(
4         'Please introduce Beijing in detail', stream=True
5 ):
6     print(response[response_printed:], end='')
7     response_printed = len(response)
File /usr/local/lib/python3.10/site-packages/EasyDel/serve/jax_serve.py:742, in JAXServer.process(self, string, greedy, max_new_tokens, **kwargs)
736 for _ in range((max_new_tokens or self.config.max_new_tokens) // self.config.max_stream_tokens):
737     inputs_to_gen = dict(
738         params=self.params,
739         input_ids=input_ids,
740         attention_mask=attention_mask
741     )
--> 742     predicted_token = self.greedy_generate(**inputs_to_gen) if greedy else self.generate(**inputs_to_gen)
744     num_generated_tokens += predicted_token.shape[-1]
745     plus_attn_mask = jnp.ones((len(attention_mask), self.config.max_stream_tokens), dtype=jnp.int32)
File /usr/local/lib/python3.10/site-packages/EasyDel/serve/jax_serve.py:316, in JAXServer.generate(self, params, input_ids, attention_mask)
314 else:
315     with self.mesh:
--> 316         return self.generate_function(
317             params, input_ids, attention_mask
318         )
[... skipping hidden 12 frame]

File /usr/local/lib/python3.10/site-packages/EasyDel/serve/jax_serve.py:253, in JAXServer.configure_generate_functions..generate(parameters, input_ids, attention_mask)
251 input_ids = with_sharding_constraint(input_ids, self.config.generation_ps)
252 attention_mask = with_sharding_constraint(attention_mask, self.config.generation_ps)
--> 253 predict = model.generate(
254     input_ids,
255     attention_mask=attention_mask,
256     params=parameters,
257     generation_config=GenerationConfig(
258         max_new_tokens=self.config.max_stream_tokens,
259
260         eos_token_id=tokenizer.eos_token_id,
261         pad_token_id=tokenizer.pad_token_id,
262         bos_token_id=tokenizer.bos_token_id,
263
264         temperature=self.config.temperature,
265         do_sample=True,
266         num_beams=1,
267         top_p=self.config.top_p,
268         top_k=self.config.top_k,
269     )
270 ).sequences[:, input_ids.shape[1]:]
271 return predict
File /usr/local/lib/python3.10/site-packages/transformers/generation/flax_utils.py:433, in FlaxGenerationMixin.generate(self, input_ids, generation_config, prng_key, trace, params, logits_processor, **kwargs)
431 elif generation_config.do_sample and generation_config.num_beams == 1:
432     logits_warper = self._get_logits_warper(generation_config=generation_config)
--> 433     return self._sample(
434         input_ids,
435         generation_config.max_length,
436         generation_config.pad_token_id,
437         generation_config.eos_token_id,
438         prng_key,
439         logits_warper=logits_warper,
440         logits_processor=logits_processor,
441         trace=trace,
442         params=params,
443         model_kwargs=model_kwargs,
444     )
445 elif not generation_config.do_sample and generation_config.num_beams > 1:
446     # broadcast input_ids & encoder_outputs
447     input_ids = self._expand_to_num_beams(input_ids, num_beams=generation_config.num_beams)
File /usr/local/lib/python3.10/site-packages/transformers/generation/flax_utils.py:737, in FlaxGenerationMixin._sample(self, input_ids, max_length, pad_token_id, eos_token_id, prng_key, logits_processor, logits_warper, trace, params, model_kwargs)
735 # The very first prompt often has sequence length > 1, so run outside of lax.while_loop to comply with TPU
736 if input_ids.shape[1] > 1:
--> 737     state = sample_search_body_fn(state)
739 if not trace:
740     state = self._run_loop_in_debug(sample_search_cond_fn, sample_search_body_fn, state)
File /usr/local/lib/python3.10/site-packages/transformers/generation/flax_utils.py:708, in FlaxGenerationMixin._sample..sample_search_body_fn(state)
706 """state update fn."""
707 prng_key, prng_key_next = jax.random.split(state.prng_key)
--> 708 model_outputs = model(state.running_token, params=params, **state.model_kwargs)
710 logits = model_outputs.logits[:, -1]
712 # apply min_length, ...
File /usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py:807, in FlaxLlamaPreTrainedModel.call(self, input_ids, attention_mask, position_ids, params, past_key_values, dropout_rng, train, output_attentions, output_hidden_states, return_dict, extra_embedding, add_params_field)
804 else:
805     mutable = False
--> 807 outputs = self.module.apply(
808     inputs,
809     jnp.array(input_ids, dtype="i4"),
810     jnp.array(attention_mask, dtype="i4"),
811     jnp.array(position_ids, dtype="i4"),
812     not train,
813     False,
814     output_attentions,
815     output_hidden_states,
816     return_dict,
817     extra_embedding,
818     rngs=rngs,
819     mutable=mutable,
820 )
822 if past_key_values is not None and return_dict:
823     outputs, past_key_values = outputs
[... skipping hidden 6 frame]

File /usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py:1116, in FlaxLlamaForCausalLMModule.call(self, input_ids, attention_mask, position_ids, deterministic, init_cache, output_attentions, output_hidden_states, return_dict, extra_embedding)
1111 if position_ids is None:
1112     position_ids = jnp.broadcast_to(
1113         jnp.clip(jnp.cumsum(attention_mask, axis=-1) - 1, a_min=0),
1114         (batch_size, seq_length)
1115     )
-> 1116 outputs = self.model(
1117     input_ids,
1118     attention_mask,
1119     position_ids,
1120     deterministic=deterministic,
1121     init_cache=init_cache,
1122     output_attentions=output_attentions,
1123     output_hidden_states=output_hidden_states,
1124     return_dict=return_dict,
1125     extra_embedding=extra_embedding
1126 )
1128 hidden_states = outputs[0]
1130 if self.config.tie_word_embeddings:
[... skipping hidden 2 frame]

File /usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py:1007, in FlaxLlamaModule.call(self, input_ids, attention_mask, position_ids, deterministic, inputs_embeds, init_cache, output_attentions, output_hidden_states, return_dict, extra_embedding)
985 """
986 The call function is the main function of a Flax model. It takes in input_ids, attention_mask, and position_ids
987 and returns the output of the model. The call function also has optional arguments that can be used to control
(...)
1004
1005 """
1006 if inputs_embeds is None:
-> 1007     inputs_embeds = self.embed_tokens(input_ids.astype("i4"))
1009 batch_size, sequence_length, _ = inputs_embeds.shape
1010 assert sequence_length <= self.config.max_position_embeddings, (f'Position out of range '
1011                                                                 f'(Model Support '
1012                                                                 f'{self.config.max_position_embeddings} got'
1013                                                                 f' {sequence_length})')
[... skipping hidden 5 frame]

File /usr/local/lib/python3.10/site-packages/flax/linen/linear.py:860, in Embed.setup(self)
859 def setup(self):
--> 860   self.embedding = self.param(
861       'embedding',
862       self.embedding_init,
863       (self.num_embeddings, self.features),
864       self.param_dtype,
865   )
[... skipping hidden 1 frame]

File /usr/local/lib/python3.10/site-packages/flax/core/scope.py:977, in Scope.param(self, name, init_fn, unbox, *init_args)
975 if not self.is_mutable_collection('params'):
976   if self.is_collection_empty('params'):
--> 977     raise errors.ScopeCollectionNotFound('params', name, self.path_text)
978   raise errors.ScopeParamNotFoundError(name, self.path_text)
979 value = init_fn(self.make_rng('params'), *init_args)
ScopeCollectionNotFound: Tried to access "embedding" from collection "params" in "/model/embed_tokens" but the collection is empty. (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ScopeCollectionNotFound)
Could you please show me how to use your partition_rules correctly?
Thank you for your quick response and patience every time.

@erfanzar I may found where the issue is. I didn't modify the partition_rules.
However, I changed the max_stream_tokens in JAXServerConfig from 32(from the documents' example code) to 1. the model's generation is ok. But the speed was lower. Why is that?

@erfanzar sorry, I'm not an expert in programing. I might do it wrong. I don't know where to put your partition_rules.
so I tried below code. but it return error when generating.
`server = JAXServer.load_from_params( params=params, model=model, config_model=config, add_params_field=True, tokenizer=tokenizer, verbose=False, do_memory_log=True, config=configs )
partition_rules = ( ("model/embed_tokens/embedding", PS("tp", ("fsdp", "sp"))), ("self_attn/(q_proj|k_proj|v_proj)/kernel", PS(("fsdp", "sp"), "tp")), ("self_attn/o_proj/kernel", PS("tp", ("fsdp", "sp"))),
        ("mlp/gate_proj/kernel", PS(("fsdp", "sp"), "tp")),
        ("mlp/down_proj/kernel", PS("tp", ("fsdp", "sp"))),
        ("mlp/up_proj/kernel", PS(("fsdp", "sp"), "tp")),

        ("input_layernorm/kernel", PS(None)),
        ("post_attention_layernorm/kernel", PS(None)),

        ("model/norm/kernel", PS(None)),
        ("lm_head/kernel", PS(("fsdp", "sp"), "tp")),
        ('.*', PS(None))

)
server.auto_configure(model,params,tokenizer,partition_rules) response_printed = 0 for response, tokens_used in server.process( 'Please introduce Beijing in detail', stream=True ): print(response[response_printed:], end='') response_printed = len(response) `
ScopeCollectionNotFound Traceback (most recent call last) Cell In[4], line 3 1 print('######################') 2 response_printed = 0 ----> 3 for response, tokens_used in server.process( 4 'Please introduce Beijing in detail', stream=True 5 ): 6 print(response[response_printed:], end='') 7 response_printed = len(response)
File /usr/local/lib/python3.10/site-packages/EasyDel/serve/jax_serve.py:742, in JAXServer.process(self, string, greedy, max_new_tokens, **kwargs) 736 for _ in range((max_new_tokens or self.config.max_new_tokens) // self.config.max_stream_tokens): 737 inputs_to_gen = dict( 738 params=self.params, 739 input_ids=input_ids, 740 attention_mask=attention_mask 741 ) --> 742 predicted_token = self.greedy_generate(**inputs_to_gen) if greedy else self.generate(**inputs_to_gen) 744 num_generated_tokens += predicted_token.shape[-1] 745 plus_attn_mask = jnp.ones((len(attention_mask), self.config.max_stream_tokens), dtype=jnp.int32)
File /usr/local/lib/python3.10/site-packages/EasyDel/serve/jax_serve.py:316, in JAXServer.generate(self, params, input_ids, attention_mask) 314 else: 315 with self.mesh: --> 316 return self.generate_function( 317 params, input_ids, attention_mask 318 )
[... skipping hidden 12 frame]

File /usr/local/lib/python3.10/site-packages/EasyDel/serve/jax_serve.py:253, in JAXServer.configure_generate_functions..generate(parameters, input_ids, attention_mask) 251 input_ids = with_sharding_constraint(input_ids, self.config.generation_ps) 252 attention_mask = with_sharding_constraint(attention_mask, self.config.generation_ps) --> 253 predict = model.generate( 254 input_ids, 255 attention_mask=attention_mask, 256 params=parameters, 257 generation_config=GenerationConfig( 258 max_new_tokens=self.config.max_stream_tokens, 259 260 eos_token_id=tokenizer.eos_token_id, 261 pad_token_id=tokenizer.pad_token_id, 262 bos_token_id=tokenizer.bos_token_id, 263 264 temperature=self.config.temperature, 265 do_sample=True, 266 num_beams=1, 267 top_p=self.config.top_p, 268 top_k=self.config.top_k, 269 ) 270 ).sequences[:, input_ids.shape[1]:] 271 return predict
File /usr/local/lib/python3.10/site-packages/transformers/generation/flax_utils.py:433, in FlaxGenerationMixin.generate(self, input_ids, generation_config, prng_key, trace, params, logits_processor, **kwargs) 431 elif generation_config.do_sample and generation_config.num_beams == 1: 432 logits_warper = self._get_logits_warper(generation_config=generation_config) --> 433 return self._sample( 434 input_ids, 435 generation_config.max_length, 436 generation_config.pad_token_id, 437 generation_config.eos_token_id, 438 prng_key, 439 logits_warper=logits_warper, 440 logits_processor=logits_processor, 441 trace=trace, 442 params=params, 443 model_kwargs=model_kwargs, 444 ) 445 elif not generation_config.do_sample and generation_config.num_beams > 1: 446 # broadcast input_ids & encoder_outputs 447 input_ids = self._expand_to_num_beams(input_ids, num_beams=generation_config.num_beams)
File /usr/local/lib/python3.10/site-packages/transformers/generation/flax_utils.py:737, in FlaxGenerationMixin._sample(self, input_ids, max_length, pad_token_id, eos_token_id, prng_key, logits_processor, logits_warper, trace, params, model_kwargs) 735 # The very first prompt often has sequence length > 1, so run outside of lax.while_loop to comply with TPU 736 if input_ids.shape[1] > 1: --> 737 state = sample_search_body_fn(state) 739 if not trace: 740 state = self._run_loop_in_debug(sample_search_cond_fn, sample_search_body_fn, state)
File /usr/local/lib/python3.10/site-packages/transformers/generation/flax_utils.py:708, in FlaxGenerationMixin._sample..sample_search_body_fn(state) 706 """state update fn.""" 707 prng_key, prng_key_next = jax.random.split(state.prng_key) --> 708 model_outputs = model(state.running_token, params=params, **state.model_kwargs) 710 logits = model_outputs.logits[:, -1] 712 # apply min_length, ...
File /usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py:807, in FlaxLlamaPreTrainedModel.call(self, input_ids, attention_mask, position_ids, params, past_key_values, dropout_rng, train, output_attentions, output_hidden_states, return_dict, extra_embedding, add_params_field) 804 else: 805 mutable = False --> 807 outputs = self.module.apply( 808 inputs, 809 jnp.array(input_ids, dtype="i4"), 810 jnp.array(attention_mask, dtype="i4"), 811 jnp.array(position_ids, dtype="i4"), 812 not train, 813 False, 814 output_attentions, 815 output_hidden_states, 816 return_dict, 817 extra_embedding, 818 rngs=rngs, 819 mutable=mutable, 820 ) 822 if past_key_values is not None and return_dict: 823 outputs, past_key_values = outputs
[... skipping hidden 6 frame]

File /usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py:1116, in FlaxLlamaForCausalLMModule.call(self, input_ids, attention_mask, position_ids, deterministic, init_cache, output_attentions, output_hidden_states, return_dict, extra_embedding) 1111 if position_ids is None: 1112 position_ids = jnp.broadcast_to( 1113 jnp.clip(jnp.cumsum(attention_mask, axis=-1) - 1, a_min=0), 1114 (batch_size, seq_length) 1115 ) -> 1116 outputs = self.model( 1117 input_ids, 1118 attention_mask, 1119 position_ids, 1120 deterministic=deterministic, 1121 init_cache=init_cache, 1122 output_attentions=output_attentions, 1123 output_hidden_states=output_hidden_states, 1124 return_dict=return_dict, 1125 extra_embedding=extra_embedding 1126 ) 1128 hidden_states = outputs[0] 1130 if self.config.tie_word_embeddings:
[... skipping hidden 2 frame]

File /usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py:1007, in FlaxLlamaModule.call(self, input_ids, attention_mask, position_ids, deterministic, inputs_embeds, init_cache, output_attentions, output_hidden_states, return_dict, extra_embedding) 985 """ 986 The call function is the main function of a Flax model. It takes in input_ids, attention_mask, and position_ids 987 and returns the output of the model. The call function also has optional arguments that can be used to control (...) 1004 1005 """ 1006 if inputs_embeds is None: -> 1007 inputs_embeds = self.embed_tokens(input_ids.astype("i4")) 1009 batch_size, sequence_length, _ = inputs_embeds.shape 1010 assert sequence_length <= self.config.max_position_embeddings, (f'Position out of range ' 1011 f'(Model Support ' 1012 f'{self.config.max_position_embeddings} got' 1013 f' {sequence_length})')
[... skipping hidden 5 frame]

File /usr/local/lib/python3.10/site-packages/flax/linen/linear.py:860, in Embed.setup(self) 859 def setup(self): --> 860 self.embedding = self.param( 861 'embedding', 862 self.embedding_init, 863 (self.num_embeddings, self.features), 864 self.param_dtype, 865 )
[... skipping hidden 1 frame]

File /usr/local/lib/python3.10/site-packages/flax/core/scope.py:977, in Scope.param(self, name, init_fn, unbox, *init_args) 975 if not self.is_mutable_collection('params'): 976 if self.is_collection_empty('params'): --> 977 raise errors.ScopeCollectionNotFound('params', name, self.path_text) 978 raise errors.ScopeParamNotFoundError(name, self.path_text) 979 value = init_fn(self.make_rng('params'), *init_args)
ScopeCollectionNotFound: Tried to access "embedding" from collection "params" in "/model/embed_tokens" but the collection is empty. (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ScopeCollectionNotFound)
Could you please show me how to use your partition_rules correctly? Thank you for your quick response and patience every time.

its ok feel free to ask anything anytime again if you want <3

@erfanzar I may find out where the issue is. I didn't modify the partition_rules. However, I changed the max_stream_tokens in JAXServerConfig from 32(from the documents' example code) to 1. the model's generation is ok. But the speed was lower. Why is that?

due to having an un-normal vocab size, we used a custom partition rule which does not focus on performance and model being scale-able and fast it just focuses on accuracy I'll write a custom Partition Rule for you and test some again to give the faster one

@erfanzar Thank you. I'm sorry to disturb you again... I have another two questions that I cannot figure out by myself.

How to do inference in 4bit? I tried config.bits=4 or 8. I found the generation speed did not change. and if config.bits=4. the output is random characters.

I tried to finetune a small model to learn easydel framework. I can successfully finetune model with below code (config.use_flash_attention=False). If I set config.use_flash_attention=True, I need manually set input_shape and init_input_shape parameter. Then, when ran the code, an error shows. How can I modify my code to use flash_attention to finetune a model? I'm running it on kaggle TPU v3-8 VM.

`
def main(argv):
params, config = llama_from_pretrained('TinyLlama/TinyLlama-1.1B-Chat-v0.6',jax.devices('cpu')[0])
config.use_flash_attention = True
config.bits = 4
train_dataset = get_dataset('erfanzar/orca-lite')
train_args = TrainArguments(
    model_class=EasyDel.modules.FlaxLlamaForCausalLM,
    configs_to_init_model_class={
        'config': config,
        'dtype': get_float_dtype_by_name('bf16'),
        'param_dtype': get_float_dtype_by_name('bf16'),
        'input_shape':(1,1024)
    },
    init_input_shape=(1,1024),
    custom_rule=config.get_partition_rules(True),
    model_name='llama',
    num_train_epochs=1,
    learning_rate=1e-4,
    learning_rate_end=1e-6,
    optimizer='adamw',
    scheduler='cosine',
    weight_decay=0.01,
    total_batch_size=8,
    max_steps=100,
    do_train=True,
    do_eval=False,
    do_test=False,
    backend='tpu',
    max_length=2048,
    gradient_checkpointing='nothing_saveable',
    sharding_array=(1, -1, 1, 1),
    use_pjit_attention_force=True,
    gradient_accumulation_steps=8,
    remove_ckpt_after_load=False,

)

trainer = CausalLanguageModelTrainer(train_args,
                          dataset_train=train_dataset['train'],
                          checkpoint_path='/kaggle/working/output')
output = trainer.train(
    model_parameters=flax.core.FrozenDict({'params': params})
)

if name == "main":
app.run(main)`
Action : Sharding Passed Parameters
Model Contain :  1.100048384  Billion Parameters
0%|                                                   | 0/100 [00:00<?, ?it/s]I1227 07:13:48.322814<PHONE_NUMBER>46624 mesh_utils.py:71] Reordering mesh to physical ring order on single-tray TPU v2/v3.
Traceback (most recent call last):
File "/root/train.py", line 230, in
app.run(main)
File "/usr/local/lib/python3.10/site-packages/absl/app.py", line 308, in run
run_main(main, args)
File "/usr/local/lib/python3.10/site-packages/absl/app.py", line 254, in run_main
sys.exit(main(argv))
File "/root/train.py", line 223, in main
output = trainer.train(
File "/usr/local/lib/python3.10/site-packages/EasyDel/trainer/causal_language_model_trainer.py", line 572, in train
sharded_train_state, loss, accuracy = self.sharded_train_step_fn(sharded_train_state,
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 255, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 161, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 767, in infer_params
return common_infer_params(pjit_info_args, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 485, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 962, in _pjit_jaxpr
jaxpr, final_consts, out_type = _create_pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 348, in memoized_fun
ans = call(fun, *args)
File "/usr/local/lib/python3.10/site-packages/jax/src/pjit.py", line 915, in create_pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2203, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/EasyDel/trainer/causal_language_model_trainer.py", line 455, in fsdp_train_step
(loss, accuracy), grad = grad_fn(state.params)
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/api.py", line 745, in value_and_grad_f
ans, vjp_py, aux = _vjp(
File "/usr/local/lib/python3.10/site-packages/jax/_src/api.py", line 2254, in _vjp
out_primal, out_vjp, aux = ad.vjp(
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 142, in vjp
out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 129, in linearize
jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals)
File "/usr/local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 777, in trace_to_jaxpr_nounits
jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/EasyDel/trainer/causal_language_model_trainer.py", line 446, in calculate_loss
logits = state.apply_fn(params=params, **batch,
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 807, in call
outputs = self.module.apply(
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1911, in apply
return apply(
File "/usr/local/lib/python3.10/site-packages/flax/core/scope.py", line 1080, in wrapper
y = fn(root, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 2572, in scope_fn
return fn(module.clone(parent=scope, _deep_clone=True), *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 584, in wrapped_module_method
return self._call_wrapped_method(fun, args, kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1101, in _call_wrapped_method
y = run_fun(self, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 1116, in call
outputs = self.model(
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 584, in wrapped_module_method
return self._call_wrapped_method(fun, args, kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1101, in _call_wrapped_method
y = run_fun(self, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 1019, in call
outputs = self.layers(
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 584, in wrapped_module_method
return self._call_wrapped_method(fun, args, kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1101, in _call_wrapped_method
y = run_fun(self, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 909, in call
layer_outputs = block(
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 584, in wrapped_module_method
return self._call_wrapped_method(fun, args, kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1101, in _call_wrapped_method
y = run_fun(self, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 575, in call
attn_outputs = self.self_attn(
File "/usr/local/lib/python3.10/site-packages/flax/linen/transforms.py", line 353, in wrapped_fn
ret = trafo_fn(module_scopes, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/core/lift.py", line 270, in wrapper
y, out_variable_groups_xs_t = fn(
File "/usr/local/lib/python3.10/site-packages/flax/linen/partitioning.py", line 553, in inner
return rematted(variable_groups, rng_groups, *dyn_args)
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/ad_checkpoint.py", line 283, in fun_remat
jaxpr, consts, out_tree = trace_to_jaxpr(fun, in_tree, tuple(in_avals))
File "/usr/local/lib/python3.10/site-packages/jax/_src/ad_checkpoint.py", line 373, in _trace_to_jaxpr
jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(flat_fun, in_avals, debug)
File "/usr/local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2203, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/flax/linen/partitioning.py", line 550, in rematted
y = fn(scope, *args)
File "/usr/local/lib/python3.10/site-packages/flax/linen/transforms.py", line 345, in core_fn
res = fn(cloned, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 584, in wrapped_module_method
return self._call_wrapped_method(fun, args, kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1101, in _call_wrapped_method
y = run_fun(self, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 348, in call
attn_output = smart_flash_attention(
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/flax_modelling_utils.py", line 464, in smart_flash_attention
attn_output = fjformer.attention.jax_flash_attn_tpu.flash_attention(
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 255, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 161, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/usr/local/lib/python3.10/site-packages/jax/_src/api.py", line 325, in infer_params
return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 485, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 962, in _pjit_jaxpr
jaxpr, final_consts, out_type = _create_pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 348, in memoized_fun
ans = call(fun, *args)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 915, in _create_pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2203, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 181, in flash_attention
return _flash_attention(
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/custom_derivatives.py", line 618, in call
out_flat = custom_vjp_call_p.bind(flat_fun, flat_fwd, flat_bwd,
File "/usr/local/lib/python3.10/site-packages/jax/src/custom_derivatives.py", line 768, in bind
outs = top_trace.process_custom_vjp_call(self, fun, fwd, bwd, tracers,
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2033, in process_custom_vjp_call
fun_jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, self.main, in_avals)
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 199, in _flash_attention
return _flash_attention_impl(
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 729, in _flash_attention_impl
o, *aux = pl.pallas_call(
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 255, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 161, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/usr/local/lib/python3.10/site-packages/jax/_src/api.py", line 325, in infer_params
return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 485, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 962, in _pjit_jaxpr
jaxpr, final_consts, out_type = _create_pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 348, in memoized_fun
ans = call(fun, *args)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 915, in _create_pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2203, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/jax/_src/pallas/pallas_call.py", line 379, in wrapped
grid_mapping, jaxpr, consts, _ = _trace_to_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pallas/pallas_call.py", line 338, in _trace_to_jaxpr
jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, jaxpr_flat_avals,
File "/usr/local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2203, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 321, in _flash_attention_kernel
kernel((batch_idx, 0), q_tile_ref, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 367, in _flash_attention_kernel_single_batch
def run():
File "/usr/local/lib/python3.10/site-packages/jax/_src/pallas/utils.py", line 29, in _wrapped
f()
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 371, in run
def body(i, _):
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
TypeError: fori_loop() got an unexpected keyword argument 'unroll'
Traceback (most recent call last):
File "/root/train.py", line 230, in
app.run(main)
File "/usr/local/lib/python3.10/site-packages/absl/app.py", line 308, in run
run_main(main, args)
File "/usr/local/lib/python3.10/site-packages/absl/app.py", line 254, in run_main
sys.exit(main(argv))
File "/root/train.py", line 223, in main
output = trainer.train(
File "/usr/local/lib/python3.10/site-packages/EasyDel/trainer/causal_language_model_trainer.py", line 572, in train
sharded_train_state, loss, accuracy = self.sharded_train_step_fn(sharded_train_state,
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 255, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 161, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 767, in infer_params
return common_infer_params(pjit_info_args, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 485, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 962, in _pjit_jaxpr
jaxpr, final_consts, out_type = _create_pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 348, in memoized_fun
ans = call(fun, *args)
File "/usr/local/lib/python3.10/site-packages/jax/src/pjit.py", line 915, in create_pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2203, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/EasyDel/trainer/causal_language_model_trainer.py", line 455, in fsdp_train_step
(loss, accuracy), grad = grad_fn(state.params)
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/api.py", line 745, in value_and_grad_f
ans, vjp_py, aux = _vjp(
File "/usr/local/lib/python3.10/site-packages/jax/_src/api.py", line 2254, in _vjp
out_primal, out_vjp, aux = ad.vjp(
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 142, in vjp
out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 129, in linearize
jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals)
File "/usr/local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 777, in trace_to_jaxpr_nounits
jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/EasyDel/trainer/causal_language_model_trainer.py", line 446, in calculate_loss
logits = state.apply_fn(params=params, **batch,
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 807, in call
outputs = self.module.apply(
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1911, in apply
return apply(
File "/usr/local/lib/python3.10/site-packages/flax/core/scope.py", line 1080, in wrapper
y = fn(root, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 2572, in scope_fn
return fn(module.clone(parent=scope, _deep_clone=True), *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 584, in wrapped_module_method
return self._call_wrapped_method(fun, args, kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1101, in _call_wrapped_method
y = run_fun(self, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 1116, in call
outputs = self.model(
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 584, in wrapped_module_method
return self._call_wrapped_method(fun, args, kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1101, in _call_wrapped_method
y = run_fun(self, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 1019, in call
outputs = self.layers(
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 584, in wrapped_module_method
return self._call_wrapped_method(fun, args, kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1101, in _call_wrapped_method
y = run_fun(self, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 909, in call
layer_outputs = block(
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 584, in wrapped_module_method
return self._call_wrapped_method(fun, args, kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1101, in _call_wrapped_method
y = run_fun(self, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 575, in call
attn_outputs = self.self_attn(
File "/usr/local/lib/python3.10/site-packages/flax/linen/transforms.py", line 353, in wrapped_fn
ret = trafo_fn(module_scopes, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/core/lift.py", line 270, in wrapper
y, out_variable_groups_xs_t = fn(
File "/usr/local/lib/python3.10/site-packages/flax/linen/partitioning.py", line 553, in inner
return rematted(variable_groups, rng_groups, *dyn_args)
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/ad_checkpoint.py", line 283, in fun_remat
jaxpr, consts, out_tree = trace_to_jaxpr(fun, in_tree, tuple(in_avals))
File "/usr/local/lib/python3.10/site-packages/jax/_src/ad_checkpoint.py", line 373, in _trace_to_jaxpr
jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(flat_fun, in_avals, debug)
File "/usr/local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2203, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/flax/linen/partitioning.py", line 550, in rematted
y = fn(scope, *args)
File "/usr/local/lib/python3.10/site-packages/flax/linen/transforms.py", line 345, in core_fn
res = fn(cloned, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 584, in wrapped_module_method
return self._call_wrapped_method(fun, args, kwargs)
File "/usr/local/lib/python3.10/site-packages/flax/linen/module.py", line 1101, in _call_wrapped_method
y = run_fun(self, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 348, in call
attn_output = smart_flash_attention(
File "/usr/local/lib/python3.10/site-packages/EasyDel/modules/flax_modelling_utils.py", line 464, in smart_flash_attention
attn_output = fjformer.attention.jax_flash_attn_tpu.flash_attention(
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 255, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 161, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/usr/local/lib/python3.10/site-packages/jax/_src/api.py", line 325, in infer_params
return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 485, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 962, in _pjit_jaxpr
jaxpr, final_consts, out_type = _create_pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 348, in memoized_fun
ans = call(fun, *args)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 915, in _create_pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2203, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 181, in flash_attention
return _flash_attention(
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/custom_derivatives.py", line 618, in call
out_flat = custom_vjp_call_p.bind(flat_fun, flat_fwd, flat_bwd,
File "/usr/local/lib/python3.10/site-packages/jax/src/custom_derivatives.py", line 768, in bind
outs = top_trace.process_custom_vjp_call(self, fun, fwd, bwd, tracers,
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2033, in process_custom_vjp_call
fun_jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, self.main, in_avals)
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 199, in _flash_attention
return _flash_attention_impl(
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 729, in _flash_attention_impl
o, *aux = pl.pallas_call(
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 255, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 161, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/usr/local/lib/python3.10/site-packages/jax/_src/api.py", line 325, in infer_params
return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 485, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 962, in _pjit_jaxpr
jaxpr, final_consts, out_type = _create_pjit_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 348, in memoized_fun
ans = call(fun, *args)
File "/usr/local/lib/python3.10/site-packages/jax/_src/pjit.py", line 915, in _create_pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2203, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/jax/_src/pallas/pallas_call.py", line 379, in wrapped
grid_mapping, jaxpr, consts, _ = _trace_to_jaxpr(
File "/usr/local/lib/python3.10/site-packages/jax/_src/pallas/pallas_call.py", line 338, in _trace_to_jaxpr
jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, jaxpr_flat_avals,
File "/usr/local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2203, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/usr/local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2225, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/usr/local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 190, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 321, in _flash_attention_kernel
kernel((batch_idx, 0), q_tile_ref, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 367, in _flash_attention_kernel_single_batch
def run():
File "/usr/local/lib/python3.10/site-packages/jax/_src/pallas/utils.py", line 29, in _wrapped
f()
File "/usr/local/lib/python3.10/site-packages/fjformer/attention/jax_flash_attn_tpu.py", line 371, in run
def body(i, _):
File "/usr/local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
TypeError: fori_loop() got an unexpected keyword argument 'unroll'
wandb:
wandb: Run history:
wandb: model billion parameters ▁
wandb:
wandb: Run summary:
wandb: model billion parameters 1.10005
wandb:
wandb: You can sync this run to the cloud by running:
wandb: wandb sync /root/wandb/offline-run-20231227_071324-3idw1ie5
wandb: Find logs at: ./wandb/offline-run-20231227_071324-3idw1ie5/logs

you have to tokenize the dataset first before passing it to train Module
here's an example:
import jax.numpy
from EasyDel import (
    TrainArguments,
    CausalLanguageModelTrainer,
    AutoEasyDelModelForCausalLM,
    EasyDelOptimizers,
    EasyDelSchedulers,
    EasyDelGradientCheckPointers
)
from datasets import load_dataset
import flax
from jax import numpy as jnp
from transformers import AutoTokenizer

huggingface_repo_id_or_path = ""

model, params = AutoEasyDelModelForCausalLM.from_pretrained(huggingface_repo_id_or_path, )

max_length = 4096
tokenizer = AutoTokenizer.from_pretrained(
    huggingface_repo_id_or_path,
    trust_remote_code=True
)
configs_to_init_model_class = {
    "config": model.config,
    "dtype": jnp.bfloat16,
    "param_dtype": jnp.bfloat16,
    "input_shape": (1, 1)
}

train_arguments = TrainArguments(
    model_class=type(model),
    model_name="my_first_model_to_train_using_easydel",
    num_train_epochs=3,
    learning_rate=5e-5,
    learning_rate_end=1e-6,
    optimizer=EasyDelOptimizers.ADAMW,  # "adamw", "lion", "adafactor" are supported
    scheduler=EasyDelSchedulers.LINEAR,
    # "linear","cosine", "none" ,"warm_up_cosine" and "warm_up_linear"  are supported
    weight_decay=0.01,
    total_batch_size=64,
    max_steps=None,  # None to let trainer Decide
    do_train=True,
    do_eval=False,  # it's optional but supported
    backend="tpu",  # default backed is set to cpu, so you must define you want to use tpu cpu or gpu
    max_length=max_length,  # Note that you have to change this in the model config too
    gradient_checkpointing=EasyDelGradientCheckPointers.NOTHING_SAVEABLE,
    sharding_array=(1, -1, 1, 1),  # the way to shard model across gpu,cpu or TPUs using sharding array (1, -1, 1, 1)
    # everything training will be in fully FSDP automatic and share data between devices
    use_pjit_attention_force=False,
    remove_ckpt_after_load=True,
    gradient_accumulation_steps=8,
    loss_remat="",
    dtype=jnp.bfloat16
)

def ultra_chat_prompting_process(
        data_chunk
):
    user_part = [
        chunk["content"] for chunk in data_chunk["messages"] if chunk["role"] == "user"
    ]
    assistant_part = [
        chunk["content"] for chunk in data_chunk["messages"] if chunk["role"] == "assistant"
    ]

    prompt = ""

    for uc, ac in zip(user_part, assistant_part):
        prompt += f"<|user|>\n{uc}</s>\n<|assistant|>\n{ac}</s>\n"

    return {"prompt": prompt}

tokenization_process = lambda data_chunk: tokenizer(
    data_chunk["prompt"],
    add_special_tokens=False,
    max_length=max_length,
    padding="max_length"
)

dataset = load_dataset("HuggingFaceH4/ultrachat_200k")
dataset_train = dataset["train_gen"].map(ultra_chat_prompting_process, num_proc=12)
dataset_train = dataset_train.map(
    tokenization_process,
    num_proc=12,
    remove_columns=dataset_train.column_names
)

# you can do the same for evaluation process dataset

trainer = CausalLanguageModelTrainer(
    train_arguments,
    dataset,
    checkpoint_path=None
)

output = trainer.train(flax.core.FrozenDict({"params": params}))
print(f"Hey ! , here's where your model saved {output.last_save_file_name}")


Change storagesystem.fill to always grant an item to orgroups when probability is equal to 1.0

About the PR
This is an attempt to change storage system to actually always spawn an item for an orgroup if the probabilities equal 1.0
It should also go through them in a random order.
(This was totally a stack overflow answer)
It does include another loop to handle the orgroup items though they're not and I can't imagine would be used to an massive extent anyhow, if there's a better way to do it I would love to hear it though.
This would fix #6834
Screenshots

Changelog

:cl:

fix: Fired delivery boy who has been skimming pizza from the emergency pizza crates

Ok, made a change to get the cumulativeprobability when grouping or groups and then use that to throw the dice roll which should mean that it can use any mix of probabilities.
Using the values on the right I was still able to spawn the entities everytime on the left:

OK, so I realized my mistake is a misunderstanding of how structs work, now I just want to make sure if I should be using a class or re-assigning that dictionary value to a new struct

Made feedback changes


Error, client state: -2

Getting an "Error, client state: -2" error. I'm unsure where to start looking for the root cause, as the CloudIoTCoreMQTTClient::mqttConnect() function doesn't make it clear where 'state' is set.
I think it may be related to a JWT auth issue, as when I turn on debug, get the generated JWT, and run it through a debugger (https://jwt.io/#debugger) the signature is invalid with my private key.
Is there any more verbose debug logging that can be used to find out what the cause of an abnormal client state could be?

All of the errors in CloudIoTCoreMQTTClient are coming from the PubSub MQTT client that it wraps. Check the State conditions for its low-level message explanations.  You may want to try one of the HTTP samples to get better information about what's going on but based on the docs, you should turn on debug logging and make sure that your client ID is correctly getting specified.

Thanks Gus. -2 is a connection failed.
Is it worth adding that State error table to the readme?

That's a good idea, a link also could be included in the debug log.

We may also move to lwMQTT, the examples in the master branch feel a bit easier to follow than the wrapper.


New Haxmedia bidder adapter

Matching code PR: https://github.com/prebid/Prebid.js/pull/6001

@rupam1999-web @ncolletti Can someone, please, merge the docs to main repo. Adapter was already merged several weeks ago

@ncolletti  Does that PR will be merge in nearest release or can be merged earlier ?


How to output a image file in css background-image to another folder?

The question likes https://github.com/webpack/file-loader/issues/32
It's seems no way to resolve this...
Thanks.

it is not a bug, you can add a query parameter name=xxx like this
    {
                test: /\.(png|jpg)$/,
                loader: 'url-loader?limit=100000&name=./imgs/[hash].[ext]'
    }

see loader-utils source code

thanks very much


Update Mapping Helper for new endpoint

Mostly refactored to use the HG util so any internal changes to API url won't break again.

Found a race bug. Need to fix.


Editor: Spanish version with English text

Godot version:
v3.0.stable.official
OS/device including version:
Ubuntu 17.10
Issue description:
The contextual help into "Crear Nuevo Node" (Create new node) window is in English, when I use Spanish version (es_AR)
I am using Spanish version (es_AR). When I open "Crear Nuevo Node" (Create new node) window and see the contextual help, its appear in English.
Steps to reproduce:
1- Open Godot Editor
2- Select es_AR language
3- Open / Create one scene
4- Push "Add a new node" button
5- Point some node, the help is in English.
Minimal reproduction project:

That's normal, that happens because the translation is not complete. This is not really a technical issue so I close it.
If you want to contribute to the spanish documentation you can use the weblate. :)

Well, the spanish translation seems quite complete actually. But the translation is currently only for the editor.
These particular tooltips are taken from the documentation, which is not translated at all currently. The focus for now is to complete the english documentation first and start translating it when it's in a good state.

These particular tooltips are taken from the documentation

Oh, I did not know that. Thanks :)


EventListener.secureConnectStart() is suspend in case of setting the phone time to future

I am using okhttp-3.14.7 library.
Case：Set the phone time to future
The following method of EventListener has been called，but no more method was called any more.
callStart
dnsStart
dnsEnd
connectStart
secureConnectStart

Can you post your EventListener implementation? SSL connections will fail if time is wrong, so almost certainly it's not hanging or suspending, just failing and throwing an exception.
Please reopen if you can provide a reproducible test case that shows it hanging when the time is changed on an emulator.