Gemma 3 model type incorrect

[deveworld]

Describe the bug
AssertionError: model type gemma3 is not defined. (upper task TaskType.CAUSAL_LM)

This bug cause task_in have only 'gemma3_text'

'gemma3_text': ModuleRegistration(module=<class 'easydel.modules.gemma3.modeling_gemma3_flax.Gemma3ForCausalLM'>, config=<class 'easydel.modules.gemma3.gemma3_configuration.Gemma3TextConfig'>, embedding_layer_names=None, layernorm_names=None [data_fields=('module', 'config', 'embedding_layer_names', 'layernorm_names'), meta_fields=()])

I think have to add 'gemma3' or replace to.

To Reproduce
Try to finetune using gemma3 model (I used SFT)

[erfanzar]

hi @deveworld sorry for being late

can you use the model with huggingface with the same method?

[deveworld]

python -m easydel.scripts.finetune.sft \
  --repo_id google/gemma-3-4b-pt \
  ...

I used above script (from https://easydel.readthedocs.io/en/latest/trc-welcome.html#supervised-fine-tuning-sft)
And AssertionError: model type gemma3 is not defined. (upper task TaskType.CAUSAL_LM) occured.

trl sft --model_name_or_path google/gemma-3-4b-pt --dataset_name stanfordnlp/imdb --output_dir gemma3-test

But it works properly in huggingface trl sft.

I think this problem can be fixed easily with adding gemma3 to task_in.

P.S.
I tested it, and confirm works perfectly.

[erfanzar]

Hi @deveworld,
Thanks for bringing this to our attention!
Regarding gemma3, you've identified a key point: as a causal language model, it shouldn't be processed using logic intended for image-text models.
The behavior you're likely seeing stems from how libraries like TRL determine the appropriate model class. They check the architectures field in the model's configuration file. If any of the listed architectures match known image-text types (like those in MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES), TRL specifically loads it using AutoModelForImageTextToText. If not, it defaults to AutoModelForCausalLM.

Here's the relevant logic from TRL:

# TRL Logic for Model Loading based on Architecture
config = AutoConfig.from_pretrained(model_args.model_name_or_path)
valid_image_text_architectures = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values()
>
# Check if any architecture listed in the config matches known image-text architectures
if config.architectures and any(arch in valid_image_text_architectures for arch in config.architectures):
    # If yes, use the image-text loader
    from transformers import AutoModelForImageTextToText
    model_kwargs.pop("use_cache", None) # Specific handling for image models
    model = AutoModelForImageTextToText.from_pretrained(model_args.model_name_or_path, **model_kwargs)
else:
    # Otherwise, assume it's a standard Causal LM
    model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, **model_kwargs)

Currently, it seems gemma3 might be inadvertently triggering this check in TRL.
We plan to adopt a similar, robust architecture-checking approach in Easydel. This will be included in an upcoming commit to ensure that causal models like gemma3 are correctly identified and loaded with the appropriate class.

Thanks again for the report!

[sparsh35]

This error still persists, i can reproduce it with latest main branch.

[erfanzar]

@sparsh35 can you share your code snippet?

[jasshone]

Hi! I am running into the same issue on 1.1.4 when running
model = ed.AutoEasyDeLModelForCausalLM.from_pretrained(
pretrained_model_name_or_path="unsloth/gemma-3-4b-pt",
dtype=dtype,
param_dtype=dtype,
config_kwargs=common_cfg,
trust_remote_code=True,
)