Repository: unslothai/unsloth
File: unsloth/models/llama.py (Python)

Code before:
                    train_embed_tokens = True
                else:
                    raise TypeError(
                        f"Unsloth: Module = {module} is not allowed. Only 'lm_head' and 'embed_tokens' is allowed."
                    )
            pass
        pass
        if isinstance(modules_to_save, (tuple, list)):
            modules_to_save = list(set(modules_to_save))
        pass

        vllm_engine = None
        if hasattr(model, "vllm_engine"):
            # Fast inference!
            vllm_engine = model.vllm_engine
            vllm_fast_generate = model.fast_generate
            vllm_fast_generate_batches = model.fast_generate_batches

            if modules_to_save is not None:
                raise NotImplementedError("Unsloth: Currently fast inference does not work with training embeddings or lm_head.")

            if bias != "none":
                raise NotImplementedError("Unsloth: Currently fast inference does not work with using biases for LoRA.")
        pass

        #does not get lora yet, so get name from model, not base model
        is_classification = "Classification" in str(type(model))

        arguments = dict(
            r                   = r,
            lora_alpha          = lora_alpha,
            target_modules      = final_modules,
            lora_dropout        = lora_dropout,
            bias                = bias,
            task_type           = TaskType.CAUSAL_LM if not is_classification else TaskType.SEQ_CLS,
            layers_to_transform = layers_to_transform,
            init_lora_weights   = init_lora_weights,
            loftq_config        = loftq_config,
            use_rslora          = use_rslora,
            modules_to_save     = modules_to_save,
            **kwargs,
        )
        if not SUPPORTS_LOFTQ:  del arguments["loftq_config"]
        if not SUPPORTS_RSLORA: del arguments["use_rslora"]

        _saved_temp_tokenizer = model._saved_temp_tokenizer

        lora_config = LoraConfig(**arguments)
        # First offload lm_head and embed_tokens to disk
        input_embeddings_device  = model.get_input_embeddings().weight.device
        if is_classification:

Review comment (suggestion):
```suggestion
        # Does not get lora yet, so get name from model, not base model
```

Code after:
        # Wrong - can cause reserved tokens to pop out!!
        # if train_embed_tokens or train_lm_head:
        #     fix_untrained_tokens(model, eps = 1e-16)
        # pass

        # Check modules_to_save
        if modules_to_save is not None:
            for module in modules_to_save:
                if module == "lm_head":
                    train_lm_head = True
                elif module == "embed_tokens":
                    train_embed_tokens = True
                else:
                    raise TypeError(
                        f"Unsloth: Module = {module} is not allowed. Only 'lm_head' and 'embed_tokens' is allowed."
                    )
            pass
        pass
        if isinstance(modules_to_save, (tuple, list)):
            modules_to_save = list(set(modules_to_save))
        pass

        vllm_engine = None
        if hasattr(model, "vllm_engine"):
            # Fast inference!
            vllm_engine = model.vllm_engine
            vllm_fast_generate = model.fast_generate
            vllm_fast_generate_batches = model.fast_generate_batches

            if modules_to_save is not None:
                raise NotImplementedError("Unsloth: Currently fast inference does not work with training embeddings or lm_head.")

            if bias != "none":
                raise NotImplementedError("Unsloth: Currently fast inference does not work with using biases for LoRA.")
        pass

        # Does not get lora yet, so get name from model, not base model
        is_classification = "Classification" in str(type(model))

        arguments = dict(
            r                   = r,
            lora_alpha          = lora_alpha,
            target_modules      = final_modules,
            lora_dropout        = lora_dropout,
            bias                = bias,
            task_type           = TaskType.CAUSAL_LM if not is_classification else TaskType.SEQ_CLS,
            layers_to_transform = layers_to_transform,
            init_lora_weights   = init_lora_weights,
            loftq_config        = loftq_config,
            use_rslora          = use_rslora,
            modules_to_save     = modules_to_save,