diff --git a/lib/factory_bot/definition.rb b/lib/factory_bot/definition.rb
index 668e637b..936babb1 100644
--- a/lib/factory_bot/definition.rb
+++ b/lib/factory_bot/definition.rb
@@ -52,6 +52,7 @@ def compile(klass = nil)
 
         declarations.attributes
 
+        self.klass ||= klass
         defined_traits.each do |defined_trait|
           defined_trait.klass ||= klass
           base_traits.each { |bt| bt.define_trait defined_trait }
@@ -152,7 +153,7 @@ def additional_traits
     end
 
     def trait_by_name(name)
-      trait_for(name) || Internal.trait_by_name(name)
+      trait_for(name) || Internal.trait_by_name(name, klass)
     end
 
     def trait_for(name)
diff --git a/lib/factory_bot/internal.rb b/lib/factory_bot/internal.rb
index 45afecb7..1d158aab 100644
--- a/lib/factory_bot/internal.rb
+++ b/lib/factory_bot/internal.rb
@@ -39,8 +39,8 @@ def register_trait(trait)
         trait
       end
 
-      def trait_by_name(name)
-        traits.find(name)
+      def trait_by_name(name, klass)
+        traits.find(name).tap { |t| t.klass = klass }
       end
 
       def register_sequence(sequence)
diff --git a/spec/acceptance/activesupport_instrumentation_spec.rb b/spec/acceptance/activesupport_instrumentation_spec.rb
index 39e5dbed..edc21537 100644
--- a/spec/acceptance/activesupport_instrumentation_spec.rb
+++ b/spec/acceptance/activesupport_instrumentation_spec.rb
@@ -114,14 +114,84 @@ def subscribed(callback, *args)
       FactoryBot.build(:user, :special)
     end
 
-    user_payload = tracked_payloads.detect { |payload| payload[:name] == :user }
-    expect(user_payload[:class]).to eq(User)
-    expect(user_payload[:attributes].map(&:name)).to eq([:email, :name])
-    expect(user_payload[:traits].map(&:name)).to eq(["special"])
-
-    special_payload = tracked_payloads.detect { |payload| payload[:name] == "special" }
-    expect(special_payload[:class]).to eq(User)
-    expect(special_payload[:attributes].map(&:name)).to eq([:name])
-    expect(special_payload[:traits].map(&:name)).to eq(["special"])
+    factory_payload = tracked_payloads.detect { |payload| payload[:name] == :user }
+    expect(factory_payload[:class]).to eq(User)
+    expect(factory_payload[:attributes].map(&:name)).to eq([:email, :name])
+    expect(factory_payload[:traits].map(&:name)).to eq(["special"])
+
+    trait_payload = tracked_payloads.detect { |payload| payload[:name] == "special" }
+    expect(trait_payload[:class]).to eq(User)
+    expect(trait_payload[:attributes].map(&:name)).to eq([:name])
+    expect(trait_payload[:traits].map(&:name)).to eq(["special"])
+  end
+
+  context "when factory with base traits" do
+    before do
+      define_model("Company", name: :string, email: :string)
+
+      FactoryBot.define do
+        trait :email do
+          email { "#{name}@example.com" }
+        end
+
+        factory :company, traits: [:email] do
+          name { "Charlie" }
+        end
+      end
+    end
+
+    it "builds the correct payload" do
+      tracked_payloads = []
+      callback = ->(_name, _start, _finish, _id, payload) { tracked_payloads << payload }
+
+      ActiveSupport::Notifications.subscribed(callback, "factory_bot.compile_factory") do
+        FactoryBot.build(:company)
+      end
+
+      factory_payload = tracked_payloads.detect { |payload| payload[:name] == :company }
+      expect(factory_payload[:class]).to eq(Company)
+      expect(factory_payload[:attributes].map(&:name)).to eq([:name])
+      expect(factory_payload[:traits].map(&:name)).to eq([])
+
+      trait_payload = tracked_payloads.detect { |payload| payload[:name] == "email" }
+      expect(trait_payload[:class]).to eq(Company)
+      expect(trait_payload[:attributes].map(&:name)).to eq([:email])
+      expect(trait_payload[:traits].map(&:name)).to eq([])
+    end
+  end
+
+  context "when factory with additional traits" do
+    before do
+      define_model("Company", name: :string, email: :string)
+
+      FactoryBot.define do
+        trait :email do
+          email { "#{name}@example.com" }
+        end
+
+        factory :company do
+          name { "Charlie" }
+        end
+      end
+    end
+
+    it "builds the correct payload" do
+      tracked_payloads = []
+      callback = ->(_name, _start, _finish, _id, payload) { tracked_payloads << payload }
+
+      ActiveSupport::Notifications.subscribed(callback, "factory_bot.compile_factory") do
+        FactoryBot.build(:company, :email)
+      end
+
+      factory_payload = tracked_payloads.detect { |payload| payload[:name] == :company }
+      expect(factory_payload[:class]).to eq(Company)
+      expect(factory_payload[:attributes].map(&:name)).to eq([:name])
+      expect(factory_payload[:traits].map(&:name)).to eq([])
+
+      trait_payload = tracked_payloads.detect { |payload| payload[:name] == "email" }
+      expect(trait_payload[:class]).to eq(Company)
+      expect(trait_payload[:attributes].map(&:name)).to eq([:email])
+      expect(trait_payload[:traits].map(&:name)).to eq([])
+    end
   end
 end
diff --git a/spec/factory_bot/internal_spec.rb b/spec/factory_bot/internal_spec.rb
index 0d4828b3..747886cb 100644
--- a/spec/factory_bot/internal_spec.rb
+++ b/spec/factory_bot/internal_spec.rb
@@ -20,9 +20,12 @@
   describe ".trait_by_name" do
     it "finds a previously registered trait" do
       trait = FactoryBot::Trait.new(:admin)
+      klass = instance_double("klass")
       FactoryBot::Internal.register_trait(trait)
 
-      expect(FactoryBot::Internal.trait_by_name(trait.name)).to eq trait
+      expect(trait.klass).to be_nil
+      expect(FactoryBot::Internal.trait_by_name(trait.name, klass)).to eq trait
+      expect(trait.klass).to eq klass
     end
   end