Skip to content

[Bug] flash-linear-attention 0.5.0 no longer registers model types on import fla #52

@ruhai-lin

Description

@ruhai-lin

Describe the Bug

The current code relies on the side effect of import fla to register FLA model types such as transformer.

However, with flash-linear-attention==0.5.0, top-level import fla no longer imports fla.models, so AutoConfig.register(...) is not triggered.

As a result, config files with:

"model_type": "transformer"

fail to load with:

KeyError: 'transformer'

ValueError: The checkpoint you are trying to load has model type `transformer`
but Transformers does not recognize this architecture.

This breaks both training and checkpoint conversion.

Suggested local fix:

import fla.models

instead of:

import fla

Affected places include:

flame/train.py
flame/utils/convert_dcp_to_hf.py
flame/utils/convert_hf_to_dcp.py
train.sh config parsing one-liner

Steps to Reproduce the Bug

Install flash-linear-attention==0.5.0, then run:

python - <<'PY'
import fla
from transformers.models.auto.configuration_auto import CONFIG_MAPPING

print("transformer" in CONFIG_MAPPING)
PY

Outputs:

False

Expected Behavior

Install flash-linear-attention==0.5.0, then run:

python - <<'PY'
import fla.models
from transformers.models.auto.configuration_auto import CONFIG_MAPPING

print("transformer" in CONFIG_MAPPING)
PY

Outputs:

True

Environment Information

  1. Torch: 2.12.0+cu130
  2. Triton: 3.7.0

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions